from __future__ import annotations

import logging

from app.clients.modelhub import ModelHubClient
from app.domain.model_type import engine_for_model
from app.domain.priorities import PRIORITY_OTHERS_DOWNLOADED
from app.settings import Settings
from app.storage.repositories import Repository

LOG = logging.getLogger(__name__)

SOURCE_MAP = {
    "HuggingFace": "HUGGING_FACE",
    "HUGGING_FACE": "HUGGING_FACE",
    "ModelScope": "MODEL_SCOPE",
    "MODEL_SCOPE": "MODEL_SCOPE",
}


def normalize_source(raw: str | None) -> str:
    if not raw:
        return "HUGGING_FACE"
    return SOURCE_MAP.get(raw, raw)


class DownloadSuccessCrawler:
    def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
        self.settings = settings
        self.repo = repo
        self.client = client

    def refresh(self) -> int:
        if not self.settings.enable_download_success_crawler:
            return 0
        imported = 0
        seen: set[str] = set()
        page = 1
        while True:
            data = self.client.list_success_download_tasks(page, self.settings.download_success_page_size)
            records = data.get("records") or []
            total = int(data.get("total") or 0)
            if not records:
                break

            for record in records:
                args = record.get("args") or {}
                model_id = args.get("model_id") or args.get("modelId")
                if not model_id or model_id in seen:
                    continue
                seen.add(model_id)
                source = normalize_source(args.get("source"))
                engine = engine_for_model(model_id)
                if engine != "vllm":
                    self.repo.upsert_candidate(
                        model_id=model_id,
                        source=source,
                        origin="model_download_success_page_gguf",
                        priority=PRIORITY_OTHERS_DOWNLOADED,
                        known_downloaded_by_others=True,
                        self_download_allowed=False,
                        notes=f"download task {record.get('id')} SUCCESS; GGUF/llamacpp candidate",
                    )
                else:
                    self.repo.upsert_candidate(
                        model_id=model_id,
                        source=source,
                        origin="model_download_success_page",
                        priority=PRIORITY_OTHERS_DOWNLOADED,
                        known_downloaded_by_others=True,
                        self_download_allowed=False,
                        notes=f"download task {record.get('id')} SUCCESS",
                    )
                imported += 1

            if self.settings.download_success_max_pages and page >= self.settings.download_success_max_pages:
                break
            if total and page * self.settings.download_success_page_size >= total:
                break
            if len(records) < self.settings.download_success_page_size:
                break
            page += 1

        LOG.info("download-success crawler imported/updated %s candidates", imported)
        return imported