from __future__ import annotations

import argparse

from app.scheduler.budget import Budget
from app.scheduler.planner import Planner
from app.settings import load_settings
from app.storage.db import connect
from app.storage.repositories import Repository


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=50)
    args = parser.parse_args()

    settings = load_settings(require_secrets=False)
    conn = connect(settings.db_path)
    try:
        repo = Repository(conn)
        Planner(repo, Budget(repo, settings.max_user_tasks), settings.default_gpu_alias_list).materialize_tasks()
        rows = repo.due_tasks(("pending", "ready_to_submit", "waiting_download", "queue_full_backoff"), limit=args.limit)
        for row in rows:
            print(f"#{row['id']} priority={row['priority']} status={row['status']} {row['model_id']} -> {row['gpu_type']} ({row['gpu_alias']})")
        print(f"budget used={repo.count_budgeted_tasks()} max={settings.max_user_tasks}")
    finally:
        conn.close()


if __name__ == "__main__":
    main()