from __future__ import annotations

import math
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Callable

from common import parse_datetime, read_json, utc_now, write_json
from history_stats import is_failure, is_success
from modelhub_client import ModelHubClientPool
from submission_claims import candidate_key


DEFAULT_GPU_STRATEGY_PATH = Path(".modelhub_state/gpu_strategy.json")
DEFAULT_REFRESH_SUBMISSIONS = 200
DEFAULT_RECENT_TERMINAL_WINDOW = 1000
DEFAULT_LONG_TERM_MIN_SAMPLES = 100
STRATEGY_STATE_VERSION = 1

LONG_TERM = "long_term"
ALL_SUPPORTED = "all_supported"
RECENT = "recent"
CATEGORIES = (LONG_TERM, ALL_SUPPORTED, RECENT)
CATEGORY_WEIGHTS = {LONG_TERM: 5, ALL_SUPPORTED: 3, RECENT: 2}

# The long-term half is split 50/30/20 across its three ranked GPUs.
LONG_TERM_GPU_PATTERN = (0, 1, 0, 2, 0, 1, 0, 1, 0, 2)


def _empty_category_counts() -> dict[str, int]:
    return {category: 0 for category in CATEGORIES}


def _gpu_name(record: dict[str, Any]) -> str:
    return str(record.get("gpuType") or record.get("targetGpu") or "").strip()


def _terminal_outcome(record: dict[str, Any]) -> str | None:
    try:
        if is_success(record):
            return "success"
        if is_failure(record):
            return "failure"
    except TypeError:
        # Be tolerant of APIs that serialize verifyResult as a string.
        copied = dict(record)
        try:
            copied["verifyResult"] = float(record.get("verifyResult"))
        except (TypeError, ValueError):
            copied["verifyResult"] = None
        if is_success(copied):
            return "success"
        if is_failure(copied):
            return "failure"
    return None


def _task_sort_key(record: dict[str, Any]) -> tuple[float, str]:
    timestamp = (
        parse_datetime(record.get("updateTime"))
        or parse_datetime(record.get("createTime"))
        or parse_datetime(record.get("submitTime"))
    )
    return (timestamp.timestamp() if timestamp else 0.0, str(record.get("taskId") or ""))


def wilson_lower_bound(successes: int, total: int, *, z: float = 1.96) -> float:
    if total <= 0:
        return 0.0
    probability = successes / total
    z_squared = z * z
    denominator = 1.0 + z_squared / total
    centre = probability + z_squared / (2.0 * total)
    margin = z * math.sqrt(
        (probability * (1.0 - probability) / total) + z_squared / (4.0 * total * total)
    )
    return max(0.0, (centre - margin) / denominator)


def _summarize_gpu_records(
    records: list[dict[str, Any]],
    supported_gpus: list[str],
) -> dict[str, dict[str, Any]]:
    supported = set(supported_gpus)
    counts: dict[str, dict[str, int]] = {
        gpu: {"success": 0, "failure": 0}
        for gpu in supported_gpus
    }
    for record in records:
        gpu = _gpu_name(record)
        if gpu not in supported:
            continue
        outcome = _terminal_outcome(record)
        if outcome is not None:
            counts[gpu][outcome] += 1

    summaries: dict[str, dict[str, Any]] = {}
    for gpu in supported_gpus:
        success = counts[gpu]["success"]
        failure = counts[gpu]["failure"]
        terminal = success + failure
        summaries[gpu] = {
            "gpu": gpu,
            "success": success,
            "failure": failure,
            "terminal": terminal,
            "successRate": success / terminal if terminal else 0.0,
            "wilsonLowerBound": wilson_lower_bound(success, terminal),
        }
    return summaries


def _rank_gpu_summaries(summaries: dict[str, dict[str, Any]]) -> list[dict[str, Any]]:
    return sorted(
        summaries.values(),
        key=lambda item: (
            -float(item["wilsonLowerBound"]),
            -float(item["successRate"]),
            -int(item["terminal"]),
            str(item["gpu"]),
        ),
    )


def build_strategy_snapshot(
    tasks: list[dict[str, Any]],
    *,
    supported_gpus: list[str],
    generated_at: datetime | None = None,
    generation: int = 0,
    recent_terminal_window: int = DEFAULT_RECENT_TERMINAL_WINDOW,
    long_term_min_samples: int = DEFAULT_LONG_TERM_MIN_SAMPLES,
    refresh_submissions: int = DEFAULT_REFRESH_SUBMISSIONS,
) -> dict[str, Any]:
    generated_at = generated_at or utc_now()
    supported_gpus = list(dict.fromkeys(gpu for gpu in supported_gpus if gpu))
    if not supported_gpus:
        raise ValueError("At least one supported GPU is required")

    all_time = _summarize_gpu_records(tasks, supported_gpus)
    all_time_ranked = _rank_gpu_summaries(all_time)
    qualified_long = [
        item for item in all_time_ranked
        if int(item["terminal"]) >= max(1, int(long_term_min_samples))
    ]
    long_term_gpus = [str(item["gpu"]) for item in qualified_long[:3]]
    if len(long_term_gpus) < min(3, len(supported_gpus)):
        for item in all_time_ranked:
            gpu = str(item["gpu"])
            if gpu not in long_term_gpus:
                long_term_gpus.append(gpu)
            if len(long_term_gpus) >= min(3, len(supported_gpus)):
                break

    terminal_tasks = [record for record in tasks if _terminal_outcome(record) is not None]
    terminal_tasks.sort(key=_task_sort_key, reverse=True)
    recent_tasks = terminal_tasks[: max(1, int(recent_terminal_window))]
    recent = _summarize_gpu_records(recent_tasks, supported_gpus)
    recent_ranked = _rank_gpu_summaries(recent)
    recent_qualified = [item for item in recent_ranked if int(item["terminal"]) >= 20]
    if recent_qualified:
        recent_gpu = str(recent_qualified[0]["gpu"])
    elif recent_tasks and recent_ranked:
        recent_gpu = str(recent_ranked[0]["gpu"])
    else:
        recent_gpu = long_term_gpus[0]

    return {
        "version": STRATEGY_STATE_VERSION,
        "generation": max(0, int(generation)),
        "generatedAt": generated_at.isoformat(),
        "historyReady": True,
        "historyTaskCount": len(tasks),
        "supportedGpus": supported_gpus,
        "longTermGpus": long_term_gpus,
        "recentGpu": recent_gpu,
        "recentTerminalCount": len(recent_tasks),
        "refreshSubmissions": max(1, int(refresh_submissions)),
        "recentTerminalWindow": max(1, int(recent_terminal_window)),
        "longTermMinSamples": max(1, int(long_term_min_samples)),
        "acceptedSinceRefresh": 0,
        "acceptedTotal": 0,
        "acceptedByCategory": _empty_category_counts(),
        "longTermStats": all_time_ranked,
        "recentStats": recent_ranked,
        "lastRefreshError": None,
        "refreshRetryAfter": None,
    }


def choose_next_category(counts: dict[str, int]) -> str:
    normalized = {category: max(0, int(counts.get(category, 0))) for category in CATEGORIES}
    next_total = sum(normalized.values()) + 1

    def deficit(category: str) -> int:
        return CATEGORY_WEIGHTS[category] * next_total - normalized[category] * sum(CATEGORY_WEIGHTS.values())

    return max(CATEGORIES, key=lambda category: (deficit(category), -CATEGORIES.index(category)))


class GPUStrategyManager:
    def __init__(
        self,
        path: Path | str = DEFAULT_GPU_STRATEGY_PATH,
        *,
        refresh_submissions: int = DEFAULT_REFRESH_SUBMISSIONS,
        recent_terminal_window: int = DEFAULT_RECENT_TERMINAL_WINDOW,
        long_term_min_samples: int = DEFAULT_LONG_TERM_MIN_SAMPLES,
        log_fn: Callable[[str], None] | None = None,
    ) -> None:
        self.path = Path(path)
        self.refresh_submissions = max(1, int(refresh_submissions))
        self.recent_terminal_window = max(1, int(recent_terminal_window))
        self.long_term_min_samples = max(1, int(long_term_min_samples))
        self.log = log_fn or (lambda message: print(message, flush=True))
        self.state: dict[str, Any] | None = None

    def _load(self) -> dict[str, Any] | None:
        try:
            value = read_json(self.path)
        except (FileNotFoundError, ValueError):
            return None
        return value if isinstance(value, dict) else None

    def _is_compatible(self, state: dict[str, Any] | None, supported_gpus: list[str]) -> bool:
        if not state or int(state.get("version") or 0) != STRATEGY_STATE_VERSION:
            return False
        return list(state.get("supportedGpus") or []) == supported_gpus

    def _refresh_due(self, state: dict[str, Any] | None, supported_gpus: list[str], now: datetime) -> tuple[bool, str]:
        if not self._is_compatible(state, supported_gpus):
            return True, "initial_or_gpu_catalog_changed"
        retry_after = parse_datetime(state.get("refreshRetryAfter")) if state is not None else None
        if state is not None and state.get("lastRefreshError") and retry_after is not None and retry_after > now:
            return False, "refresh_error_backoff"
        if not bool(state.get("historyReady", False)):
            return (retry_after is None or retry_after <= now), "history_not_ready"
        accepted = int(state.get("acceptedSinceRefresh") or 0)
        refresh_every = max(1, int(state.get("refreshSubmissions") or self.refresh_submissions))
        return accepted >= refresh_every, "accepted_submission_threshold"

    @staticmethod
    def _load_platform_history(client: Any) -> list[dict[str, Any]]:
        if isinstance(client, ModelHubClientPool):
            by_account: dict[int, list[dict[str, Any]]] = {}
            errors: list[int] = []
            with ThreadPoolExecutor(max_workers=min(len(client.clients), 12)) as executor:
                futures = {
                    executor.submit(account_client.list_tasks, page_size=100, only_mine=True): index
                    for index, account_client in enumerate(client.clients, start=1)
                }
                for future in as_completed(futures):
                    index = futures[future]
                    try:
                        by_account[index] = future.result()
                    except Exception:
                        errors.append(index)
            if errors:
                joined = ",".join(str(index) for index in sorted(errors))
                raise RuntimeError(f"ModelHub history fetch failed for account indexes: {joined}")
            tasks = [task for index in sorted(by_account) for task in by_account[index]]
        else:
            tasks = client.list_tasks(page_size=100, only_mine=True)

        deduped: list[dict[str, Any]] = []
        seen_task_ids: set[str] = set()
        for task in tasks:
            if not isinstance(task, dict):
                continue
            task_id = str(task.get("taskId")) if task.get("taskId") is not None else None
            if task_id and task_id in seen_task_ids:
                continue
            if task_id:
                seen_task_ids.add(task_id)
            deduped.append(task)
        return deduped

    def prepare(
        self,
        client: Any,
        *,
        supported_gpus: list[str],
        now: datetime | None = None,
    ) -> dict[str, Any]:
        now = now or utc_now()
        supported_gpus = list(dict.fromkeys(gpu for gpu in supported_gpus if gpu))
        state = self._load()
        refresh_due, reason = self._refresh_due(state, supported_gpus, now)
        if not refresh_due and state is not None:
            self.state = state
            self._log_state("loaded")
            return state

        previous_generation = int(state.get("generation", -1)) if state is not None else -1
        self.log(f"[strategy] refresh_start reason={reason} supported_gpus={len(supported_gpus)}")
        try:
            tasks = self._load_platform_history(client)
            refreshed = build_strategy_snapshot(
                tasks,
                supported_gpus=supported_gpus,
                generated_at=now,
                generation=previous_generation + 1,
                recent_terminal_window=self.recent_terminal_window,
                long_term_min_samples=self.long_term_min_samples,
                refresh_submissions=self.refresh_submissions,
            )
            refreshed["acceptedTotal"] = int((state or {}).get("acceptedTotal") or 0)
            self.state = refreshed
            write_json(self.path, refreshed)
            self._log_state("refreshed")
            return refreshed
        except Exception as exc:
            self.log(f"[strategy] refresh_error reason={type(exc).__name__}: {exc}")
            if self._is_compatible(state, supported_gpus):
                assert state is not None
                state["lastRefreshError"] = str(exc)
                state["refreshRetryAfter"] = (now + timedelta(minutes=5)).isoformat()
                self.state = state
                write_json(self.path, state)
                return state

            fallback = build_strategy_snapshot(
                [],
                supported_gpus=supported_gpus,
                generated_at=now,
                generation=0,
                recent_terminal_window=self.recent_terminal_window,
                long_term_min_samples=self.long_term_min_samples,
                refresh_submissions=self.refresh_submissions,
            )
            fallback["historyReady"] = False
            fallback["lastRefreshError"] = str(exc)
            fallback["refreshRetryAfter"] = (now + timedelta(minutes=5)).isoformat()
            self.state = fallback
            write_json(self.path, fallback)
            self._log_state("fallback")
            return fallback

    def _log_state(self, action: str) -> None:
        if self.state is None:
            return
        counts = self.state.get("acceptedByCategory") or {}
        self.log(
            f"[strategy] {action} generation={self.state.get('generation', 0)} "
            f"accepted={self.state.get('acceptedSinceRefresh', 0)}/{self.state.get('refreshSubmissions', self.refresh_submissions)} "
            f"categories={counts.get(LONG_TERM, 0)},{counts.get(ALL_SUPPORTED, 0)},{counts.get(RECENT, 0)} "
            f"long={','.join(self.state.get('longTermGpus') or [])} recent={self.state.get('recentGpu') or 'n/a'}"
        )

    @property
    def submissions_until_refresh(self) -> int:
        if self.state is None:
            return self.refresh_submissions
        refresh_every = max(1, int(self.state.get("refreshSubmissions") or self.refresh_submissions))
        accepted = max(0, int(self.state.get("acceptedSinceRefresh") or 0))
        return max(0, refresh_every - accepted)

    def _category_gpu_order(self, category: str, occurrence: int) -> list[str]:
        assert self.state is not None
        supported = list(self.state.get("supportedGpus") or [])
        long_term = [gpu for gpu in self.state.get("longTermGpus") or [] if gpu in supported]
        recent_gpu = str(self.state.get("recentGpu") or "")

        if category == LONG_TERM and long_term:
            desired_index = LONG_TERM_GPU_PATTERN[occurrence % len(LONG_TERM_GPU_PATTERN)]
            desired = long_term[min(desired_index, len(long_term) - 1)]
            return [desired, *(gpu for gpu in long_term if gpu != desired)]
        if category == RECENT and recent_gpu in supported:
            return [recent_gpu]
        if supported:
            offset = occurrence % len(supported)
            return [*supported[offset:], *supported[:offset]]
        return []

    def order_candidates(self, candidates: list[dict[str, Any]]) -> list[dict[str, Any]]:
        if self.state is None or not candidates:
            return list(candidates)

        by_gpu: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for candidate in candidates:
            by_gpu[str(candidate.get("targetGpu") or "")].append(candidate)
        gpu_indexes: dict[str, int] = defaultdict(int)
        used: set[str] = set()
        ordered: list[dict[str, Any]] = []
        virtual_counts = {
            category: max(0, int((self.state.get("acceptedByCategory") or {}).get(category, 0)))
            for category in CATEGORIES
        }

        def take_from_gpu(gpu: str) -> dict[str, Any] | None:
            pool = by_gpu.get(gpu) or []
            index = gpu_indexes[gpu]
            while index < len(pool):
                candidate = pool[index]
                index += 1
                gpu_indexes[gpu] = index
                if candidate_key(candidate) not in used:
                    return candidate
            gpu_indexes[gpu] = index
            return None

        while len(ordered) < len(candidates):
            planned_category = choose_next_category(virtual_counts)
            selected: dict[str, Any] | None = None
            actual_category = planned_category
            categories_to_try = [planned_category, *(category for category in CATEGORIES if category != planned_category)]
            for category in categories_to_try:
                occurrence = virtual_counts[category]
                for gpu in self._category_gpu_order(category, occurrence):
                    selected = take_from_gpu(gpu)
                    if selected is not None:
                        actual_category = category
                        break
                if selected is not None:
                    break

            if selected is None:
                # Unknown/custom GPUs can only appear when callers bypass the normal resolver.
                selected = next((candidate for candidate in candidates if candidate_key(candidate) not in used), None)
                actual_category = ALL_SUPPORTED
            if selected is None:
                break

            key = candidate_key(selected)
            used.add(key)
            annotated = dict(selected)
            annotated["strategyCategory"] = actual_category
            annotated["strategyPlannedCategory"] = planned_category
            annotated["strategyGeneration"] = int(self.state.get("generation") or 0)
            ordered.append(annotated)
            virtual_counts[actual_category] += 1

        return ordered

    def record_accepted(self, candidates: list[dict[str, Any]]) -> dict[str, Any] | None:
        if not candidates:
            return self.state
        state = self._load() or self.state
        if state is None:
            return None

        category_counts = {
            category: max(0, int((state.get("acceptedByCategory") or {}).get(category, 0)))
            for category in CATEGORIES
        }
        for candidate in candidates:
            category = str(candidate.get("strategyCategory") or ALL_SUPPORTED)
            if category not in category_counts:
                category = ALL_SUPPORTED
            category_counts[category] += 1

        accepted_count = len(candidates)
        state["acceptedSinceRefresh"] = int(state.get("acceptedSinceRefresh") or 0) + accepted_count
        state["acceptedTotal"] = int(state.get("acceptedTotal") or 0) + accepted_count
        state["acceptedByCategory"] = category_counts
        self.state = state
        write_json(self.path, state)
        self._log_state("progress")
        return state

    def summary(self) -> dict[str, Any]:
        if self.state is None:
            return {"enabled": False}
        return {
            "enabled": True,
            "statePath": str(self.path),
            "generation": int(self.state.get("generation") or 0),
            "generatedAt": self.state.get("generatedAt"),
            "historyReady": bool(self.state.get("historyReady", False)),
            "historyTaskCount": int(self.state.get("historyTaskCount") or 0),
            "acceptedSinceRefresh": int(self.state.get("acceptedSinceRefresh") or 0),
            "refreshSubmissions": int(self.state.get("refreshSubmissions") or self.refresh_submissions),
            "acceptedByCategory": dict(self.state.get("acceptedByCategory") or {}),
            "longTermGpus": list(self.state.get("longTermGpus") or []),
            "recentGpu": self.state.get("recentGpu"),
            "recentTerminalCount": int(self.state.get("recentTerminalCount") or 0),
            "refreshDue": self.submissions_until_refresh <= 0,
        }