from app.clients.modelhub import ModelHubClient
from app.storage.repositories import Repository, future_seconds


class Downloader:
    def __init__(self, repo: Repository, client: ModelHubClient, max_self_downloads: int, poll_interval_seconds: int):
        self.repo = repo
        self.client = client
        self.max_self_downloads = max_self_downloads
        self.poll_interval_seconds = poll_interval_seconds

    def poll_active_downloads(self) -> int:
        changed = 0
        for download in self.repo.active_downloads_due(limit=20):
            status = self.client.get_download_status(download["model_id"])
            active = status in {"WAITING", "RUNNING", "QUERY_FAILED"}
            self.repo.upsert_download(
                download["model_id"],
                download["source"],
                status,
                "self",
                active_self_slot=active,
                next_poll_at=future_seconds(self.poll_interval_seconds) if active else None,
            )
            changed += 1
        return changed

    def maybe_start_self_downloads(self, limit: int = 5) -> int:
        started = 0
        for task in self.repo.due_tasks(("waiting_download",), limit=limit):
            if self.repo.count_active_self_downloads() >= self.max_self_downloads:
                self.repo.mark_task_status(task["id"], "waiting_download", next_attempt_at=future_seconds(300))
                continue
            sync_result = self.client.sync_from_huggingface(task["model_id"])
            if sync_result.result != "success":
                self.repo.mark_task_status(task["id"], "waiting_download", sync_result.code, sync_result.message, future_seconds(600), increment_attempt=True)
                continue
            dl = self.client.create_download_task(task["model_id"], task["source"] or "HUGGING_FACE")
            if dl.result == "created":
                self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "WAITING", "self", True, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
                self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
                started += 1
            elif dl.result == "already_downloaded":
                self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "SUCCESS", "others", False, dl.code, dl.message)
                self.repo.mark_task_status(task["id"], "ready_to_submit", dl.code, dl.message)
            elif dl.result == "running_elsewhere":
                self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "RUNNING", "others", False, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
                self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
            elif dl.result == "download_pool_full":
                self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(300), increment_attempt=True)
            else:
                self.repo.mark_task_status(task["id"], "failed", dl.code, dl.message, increment_attempt=True)
        return started