from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterable

from candidate_preflight import CandidatePreflightAdvisor, MODEL_LOAD_OVERHEAD
from common import write_json
from hf_discovery import HuggingFaceDiscovery, inspect_repo_tree
from modelhub_client import ModelHubClient, ModelHubClientPool
from runner_common import DEFAULT_KEY_PATH, ensure_tokens


ACTIVE_FILTER_STATUSES = ("waiting", "running")
DEFAULT_STOP_BATCH_SIZE = 50


@dataclass(frozen=True)
class OwnedTask:
    account_index: int
    task_id: int
    model_id: str
    gpu_type: str
    status: str


def _normalize_status(value: Any) -> str:
    return str(value or "").strip().lower()


def _parse_owned_task(account_index: int, record: dict[str, Any]) -> OwnedTask | None:
    task_id = record.get("taskId")
    model_id = str(record.get("modelId") or "").strip()
    gpu_type = str(record.get("gpuType") or "").strip()
    status = _normalize_status(record.get("status"))
    if not model_id or not gpu_type or status not in ACTIVE_FILTER_STATUSES:
        return None
    try:
        numeric_task_id = int(str(task_id).strip())
    except (TypeError, ValueError):
        return None
    if numeric_task_id <= 0:
        return None
    return OwnedTask(
        account_index=account_index,
        task_id=numeric_task_id,
        model_id=model_id,
        gpu_type=gpu_type,
        status=status,
    )


def _fetch_status_tasks(
    client: ModelHubClient,
    *,
    account_index: int,
    status: str,
    page_size: int = 100,
) -> list[OwnedTask]:
    current = 1
    results: list[OwnedTask] = []
    while True:
        payload = client.list_tasks_page(
            current=current,
            page_size=page_size,
            only_mine=True,
            status=status,
        )
        page = payload.get("data") or {}
        records = page.get("records") or []
        if not isinstance(records, list):
            raise ValueError("ModelHub task page records are invalid")
        for record in records:
            if not isinstance(record, dict):
                continue
            task = _parse_owned_task(account_index, record)
            # The platform has silently ignored unknown status filters before.
            # Accept only records whose returned state is explicitly active.
            if task is not None:
                results.append(task)
        pages = int(page.get("pages") or 0)
        if not records or current >= pages:
            break
        current += 1
    return results


def collect_active_tasks(
    clients: list[ModelHubClient],
    *,
    read_concurrency: int = 6,
    statuses: Iterable[str] = ACTIVE_FILTER_STATUSES,
) -> tuple[list[OwnedTask], dict[int, list[str]]]:
    """Fetch active tasks from every account without sharing account-scoped reads."""
    requested_statuses = tuple(dict.fromkeys(_normalize_status(status) for status in statuses))
    tasks: list[OwnedTask] = []
    errors: dict[int, list[str]] = {}
    jobs = [(index, client, status) for index, client in enumerate(clients) for status in requested_statuses]
    workers = min(max(1, int(read_concurrency)), max(1, len(jobs)))
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(
                _fetch_status_tasks,
                client,
                account_index=index,
                status=status,
            ): (index, status)
            for index, client, status in jobs
        }
        for future in as_completed(futures):
            index, status = futures[future]
            try:
                tasks.extend(future.result())
            except Exception as exc:
                errors.setdefault(index, []).append(f"{status}: {type(exc).__name__}: {exc}")

    deduped: dict[tuple[int, int], OwnedTask] = {}
    for task in tasks:
        deduped[(task.account_index, task.task_id)] = task
    return sorted(deduped.values(), key=lambda item: (item.account_index, item.task_id)), errors


def _load_repository_sizes(
    model_ids: set[str],
    *,
    discovery: HuggingFaceDiscovery,
    read_concurrency: int,
    log: Callable[[str], None],
) -> tuple[dict[str, int], dict[str, str]]:
    sizes: dict[str, int] = {}
    errors: dict[str, str] = {}
    completed = 0
    workers = min(max(1, int(read_concurrency)), max(1, len(model_ids)))

    def inspect(model_id: str) -> tuple[str, int | None]:
        tree = discovery.list_repo_tree(model_id)
        return model_id, inspect_repo_tree(model_id, tree).repository_size_bytes

    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {executor.submit(inspect, model_id): model_id for model_id in sorted(model_ids)}
        for future in as_completed(futures):
            model_id = futures[future]
            try:
                returned_model_id, size = future.result()
                if size is None:
                    errors[returned_model_id] = "recursive_repository_size_incomplete"
                else:
                    sizes[returned_model_id] = size
            except Exception as exc:
                errors[model_id] = f"{type(exc).__name__}: {exc}"
            completed += 1
            if completed == len(model_ids) or completed % 50 == 0:
                log(
                    f"[queue-cleanup] size_scan={completed}/{len(model_ids)} "
                    f"complete={len(sizes)} unknown={len(errors)}"
                )
    return sizes, errors


def find_certain_oom_tasks(
    tasks: list[OwnedTask],
    *,
    repository_sizes: dict[str, int],
    gpu_memory_gib: dict[str, float] | None = None,
) -> tuple[list[dict[str, Any]], dict[str, int]]:
    capacities = dict(CandidatePreflightAdvisor(gpu_memory_gib=gpu_memory_gib).gpu_memory_gib)
    decisions: list[dict[str, Any]] = []
    skipped = {
        "repositorySizeUnknown": 0,
        "gpuCapacityUnknown": 0,
        "fitsKnownCapacity": 0,
    }
    for task in tasks:
        size_bytes = repository_sizes.get(task.model_id)
        if size_bytes is None:
            skipped["repositorySizeUnknown"] += 1
            continue
        capacity_gib = capacities.get(task.gpu_type)
        if capacity_gib is None:
            skipped["gpuCapacityUnknown"] += 1
            continue
        repository_gib = size_bytes / (1024**3)
        required_gib = repository_gib * MODEL_LOAD_OVERHEAD
        if required_gib <= capacity_gib:
            skipped["fitsKnownCapacity"] += 1
            continue
        decisions.append(
            {
                "accountIndex": task.account_index + 1,
                "taskId": task.task_id,
                "modelId": task.model_id,
                "gpuType": task.gpu_type,
                "status": task.status,
                "repositorySizeGiB": round(repository_gib, 3),
                "requiredGiB": round(required_gib, 3),
                "gpuCapacityGiB": float(capacity_gib),
                "reason": "certain_oom_repository_size_exceeds_gpu_capacity",
            }
        )
    return decisions, skipped


def _chunks(values: list[int], size: int) -> Iterable[list[int]]:
    for offset in range(0, len(values), size):
        yield values[offset : offset + size]


def cleanup_certain_oom_tasks(
    modelhub: ModelHubClientPool,
    discovery: HuggingFaceDiscovery,
    *,
    dry_run: bool = False,
    read_concurrency: int = 6,
    stop_batch_size: int = DEFAULT_STOP_BATCH_SIZE,
    gpu_memory_gib: dict[str, float] | None = None,
    log: Callable[[str], None] = print,
) -> dict[str, Any]:
    """Stop only active tasks that cannot fit the selected GPU by known capacity."""
    clients = list(modelhub.clients)
    tasks, listing_errors = collect_active_tasks(clients, read_concurrency=read_concurrency)
    log(
        f"[queue-cleanup] active_scanned={len(tasks)} accounts={len(clients)} "
        f"listing_errors={sum(len(items) for items in listing_errors.values())}"
    )

    model_ids = {task.model_id for task in tasks}
    repository_sizes, size_errors = _load_repository_sizes(
        model_ids,
        discovery=discovery,
        read_concurrency=read_concurrency,
        log=log,
    )
    decisions, skipped = find_certain_oom_tasks(
        tasks,
        repository_sizes=repository_sizes,
        gpu_memory_gib=gpu_memory_gib,
    )
    log(
        f"[queue-cleanup] certain_oom={len(decisions)} "
        f"fits={skipped['fitsKnownCapacity']} size_unknown={skipped['repositorySizeUnknown']} "
        f"gpu_unknown={skipped['gpuCapacityUnknown']} dry_run={str(bool(dry_run)).lower()}"
    )

    cancelled: list[dict[str, Any]] = []
    disappeared: list[dict[str, Any]] = []
    stop_errors: list[dict[str, Any]] = []
    if decisions and not dry_run:
        # Re-read each account immediately before mutation. If any active-state
        # query fails for that account, fail closed and do not terminate its tasks.
        refreshed_tasks, refresh_errors = collect_active_tasks(clients, read_concurrency=read_concurrency)
        active_ids_by_account: dict[int, set[int]] = {}
        for task in refreshed_tasks:
            active_ids_by_account.setdefault(task.account_index, set()).add(task.task_id)

        for account_index, details in sorted(refresh_errors.items()):
            stop_errors.append(
                {
                    "accountIndex": account_index + 1,
                    "error": "active_task_recheck_failed",
                    "details": details,
                }
            )

        by_account: dict[int, list[dict[str, Any]]] = {}
        for decision in decisions:
            account_index = int(decision["accountIndex"]) - 1
            if account_index in refresh_errors:
                continue
            if int(decision["taskId"]) not in active_ids_by_account.get(account_index, set()):
                disappeared.append(decision)
                continue
            by_account.setdefault(account_index, []).append(decision)

        batch_size = max(1, min(100, int(stop_batch_size)))
        stop_failed = False
        for account_index in sorted(by_account):
            if stop_failed:
                break
            decisions_by_id = {int(item["taskId"]): item for item in by_account[account_index]}
            task_ids = sorted(decisions_by_id)
            for batch in _chunks(task_ids, batch_size):
                try:
                    clients[account_index].stop_tasks(batch)
                except Exception as exc:
                    stop_errors.append(
                        {
                            "accountIndex": account_index + 1,
                            "taskIds": batch,
                            "error": f"{type(exc).__name__}: {exc}",
                        }
                    )
                    stop_failed = True
                    break
                cancelled.extend(decisions_by_id[task_id] for task_id in batch)
                log(
                    f"[queue-cleanup] account={account_index + 1:02d} "
                    f"cancelled_batch={len(batch)} cancelled_total={len(cancelled)}"
                )

        if cancelled and hasattr(modelhub, "refresh_active_counts"):
            modelhub.refresh_active_counts()

    return {
        "dryRun": bool(dry_run),
        "accounts": len(clients),
        "activeScanned": len(tasks),
        "uniqueModels": len(model_ids),
        "repositorySizesComplete": len(repository_sizes),
        "repositorySizeErrors": size_errors,
        "listingErrors": {str(index + 1): values for index, values in listing_errors.items()},
        "certainOomCount": len(decisions),
        "certainOomTasks": decisions,
        "skipped": skipped,
        "cancelledCount": len(cancelled),
        "cancelledTasks": cancelled,
        "noLongerActiveCount": len(disappeared),
        "stopErrors": stop_errors,
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Safely stop active ModelHub tasks that are certain to run out of GPU memory.")
    parser.add_argument("--execute", action="store_true", help="Actually terminate selected tasks; otherwise only print a preview")
    parser.add_argument("--read-concurrency", type=int, default=6)
    parser.add_argument("--stop-batch-size", type=int, default=DEFAULT_STOP_BATCH_SIZE)
    parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH))
    parser.add_argument("--modelhub-base-url", default="https://modelhub.org.cn")
    parser.add_argument("--hf-base-url", default="https://modelscope.cn")
    parser.add_argument("--modelhub-token", default=None)
    parser.add_argument("--hf-token", default=None)
    parser.add_argument("--modelscope-token", default=None)
    parser.add_argument("--report-path", default=".modelhub_state/queue_cleanup_latest.json")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    ensure_tokens(args)
    tokens = list(getattr(args, "modelhub_tokens", None) or [args.modelhub_token])
    clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in tokens]
    pool = ModelHubClientPool(clients, capacity_state_path=None)
    discovery = HuggingFaceDiscovery(base_url=args.hf_base_url)
    summary = cleanup_certain_oom_tasks(
        pool,
        discovery,
        dry_run=not args.execute,
        read_concurrency=args.read_concurrency,
        stop_batch_size=args.stop_batch_size,
    )
    report_path = Path(args.report_path)
    write_json(report_path, summary)
    print(
        f"[queue-cleanup] finished certain_oom={summary['certainOomCount']} "
        f"cancelled={summary['cancelledCount']} stop_errors={len(summary['stopErrors'])} "
        f"report={report_path}",
        flush=True,
    )
    return 1 if summary["stopErrors"] else 0


if __name__ == "__main__":
    raise SystemExit(main())