from __future__ import annotations

import logging

from app.clients.modelhub import ModelHubClient
from app.domain.machines import machine_to_gpu_alias, parse_machine_names
from app.domain.model_type import engine_for_model
from app.domain.priorities import PRIORITY_BOUNTY
from app.settings import Settings
from app.storage.repositories import Repository

LOG = logging.getLogger(__name__)


def _model_task_name(model: dict) -> str:
    task_info = model.get("taskLevelsInfo") or {}
    return model.get("taskLevelChineseName") or task_info.get("taskLevelChineseName") or ""


def _adapted_machine_names(model: dict) -> set[str]:
    machines = model.get("machines") or []
    return {m.get("machineName") for m in machines if m.get("machineName")}


class NotAdaptedCrawler:
    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_not_adapted_crawler:
            return 0
        machine_names = parse_machine_names(self.settings.target_machine_names)
        if not machine_names:
            LOG.info("not-adapted crawler skipped: TARGET_MACHINE_NAMES is empty")
            return 0

        machine_alias_pairs: list[tuple[str, str]] = []
        for machine_name in machine_names:
            alias = machine_to_gpu_alias(machine_name)
            if not alias:
                LOG.warning("unknown target machine name, skipped: %s", machine_name)
                continue
            machine_alias_pairs.append((machine_name, alias))
        if not machine_alias_pairs:
            return 0

        imported = 0
        page = 1
        while True:
            data = self.client.list_modelhub_models(page, self.settings.modelhub_page_size, "")
            records = data.get("records") or []
            total = int(data.get("total") or 0)
            if not records:
                break

            for model in records:
                model_id = model.get("modelId")
                if not model_id:
                    continue
                task_name = _model_task_name(model)
                if self.settings.target_task_level and self.settings.target_task_level not in task_name:
                    continue

                adapted = _adapted_machine_names(model)
                engine = engine_for_model(model_id)
                for machine_name, alias in machine_alias_pairs:
                    if machine_name in adapted:
                        continue
                    if engine != "vllm":
                        self.repo.upsert_candidate(
                            model_id=model_id,
                            source="HUGGING_FACE",
                            origin="bounty_target_not_adapted_gguf",
                            priority=PRIORITY_BOUNTY,
                            is_bounty=True,
                            is_promote=True,
                            known_downloaded_by_others=True,
                            self_download_allowed=False,
                            target_gpu_aliases=alias,
                            notes=f"GGUF/llamacpp candidate for {machine_name}; vLLM phase skips submit",
                        )
                        imported += 1
                        continue
                    self.repo.upsert_candidate(
                        model_id=model_id,
                        source="HUGGING_FACE",
                        origin="bounty_target_not_adapted",
                        priority=PRIORITY_BOUNTY,
                        is_bounty=True,
                        is_promote=True,
                        known_downloaded_by_others=True,
                        self_download_allowed=False,
                        target_gpu_aliases=alias,
                        notes=f"not adapted on {machine_name}",
                    )
                    imported += 1

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

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