"""Model discovery and inference smoke-test logic for the StockAnalysis Space.

The module deliberately keeps inference cheap: it resolves the most-trending
conversational models that currently have at least one live provider, then sends
one tiny chat request per model. All tunables are applied through
``configure()``, which ``app.py`` calls with its SETTINGS block.
"""

from __future__ import annotations

import json
import os
import re
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

import requests
from huggingface_hub import InferenceClient, InferenceTimeoutError
from huggingface_hub.errors import HfHubHTTPError

HUB_MODELS_API = "https://huggingface.co/api/models"
WHOAMI_API = "https://huggingface.co/api/whoami-v2"

RESULTS_ROOT = Path("/data/hf-inference-smoke")
RUNS_DIR = RESULTS_ROOT / "runs"
MODELS_DIR = RESULTS_ROOT / "models"

DEFAULT_MODEL_LIMIT = 30
DEFAULT_PAGE_SIZE = 100
DEFAULT_MAX_SCAN = 500

RUN_ID_RE = re.compile(r"^[0-9TZ]+-[0-9a-f]{6}$")

PROMPT = "ping"
MAX_TOKENS = 8
DEFAULT_TIMEOUT = 60.0
DEFAULT_CONCURRENCY = 3
ATTEMPT_POLICIES = ("cheapest", "fastest")


def configure(
    *,
    results_root: str | Path | None = None,
    model_limit: int | None = None,
    concurrency: int | None = None,
    max_tokens: int | None = None,
    request_timeout: float | None = None,
    page_size: int | None = None,
    max_scan: int | None = None,
) -> dict[str, Any]:
    """Apply the settings declared in ``app.py``.

    Every value is optional so tests can override just the storage location.
    The app calls this once at import time, which keeps all user-facing knobs
    in a single place instead of scattered environment variables.
    """
    global RESULTS_ROOT, RUNS_DIR, MODELS_DIR
    global DEFAULT_MODEL_LIMIT, DEFAULT_CONCURRENCY, MAX_TOKENS
    global DEFAULT_TIMEOUT, DEFAULT_PAGE_SIZE, DEFAULT_MAX_SCAN

    if results_root is not None:
        RESULTS_ROOT = Path(results_root)
        RUNS_DIR = RESULTS_ROOT / "runs"
        MODELS_DIR = RESULTS_ROOT / "models"
    if model_limit is not None:
        DEFAULT_MODEL_LIMIT = int(model_limit)
    if concurrency is not None:
        DEFAULT_CONCURRENCY = int(concurrency)
    if max_tokens is not None:
        MAX_TOKENS = int(max_tokens)
    if request_timeout is not None:
        DEFAULT_TIMEOUT = float(request_timeout)
    if page_size is not None:
        DEFAULT_PAGE_SIZE = int(page_size)
    if max_scan is not None:
        DEFAULT_MAX_SCAN = int(max_scan)

    return {
        "results_root": str(RESULTS_ROOT),
        "model_limit": DEFAULT_MODEL_LIMIT,
        "concurrency": DEFAULT_CONCURRENCY,
        "max_tokens": MAX_TOKENS,
        "request_timeout": DEFAULT_TIMEOUT,
        "page_size": DEFAULT_PAGE_SIZE,
        "max_scan": DEFAULT_MAX_SCAN,
    }


def utc_now() -> str:
    """Return an ISO-8601 UTC timestamp with a trailing Z."""
    return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")


def new_run_id() -> str:
    """Return a sortable, path-safe run identifier."""
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    return f"{stamp}-{uuid.uuid4().hex[:6]}"


def ensure_storage() -> None:
    """Create the persistent-storage directories on demand."""
    RUNS_DIR.mkdir(parents=True, exist_ok=True)
    MODELS_DIR.mkdir(parents=True, exist_ok=True)


def _write_json(path: Path, payload: Any) -> None:
    """Write JSON atomically so a crashed run cannot leave a partial file."""
    ensure_storage()
    tmp_path = path.with_suffix(path.suffix + ".tmp")
    tmp_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    tmp_path.replace(path)


def _auth_headers() -> dict[str, str]:
    token = os.getenv("HF_TOKEN")
    return {"Authorization": f"Bearer {token}"} if token else {}


def _parse_provider_mapping(raw_mapping: Any) -> tuple[list[str], dict[str, str]]:
    """Return live providers and every provider status from a mapping payload."""
    live: list[str] = []
    statuses: dict[str, str] = {}
    if not isinstance(raw_mapping, list):
        return live, statuses

    for item in raw_mapping:
        if not isinstance(item, dict):
            continue
        provider = item.get("provider")
        status = item.get("status")
        if not provider or not status:
            continue
        statuses[provider] = status
        if status == "live" and provider not in live:
            live.append(provider)
    return live, statuses


def fetch_trending_models(
    limit: int | None = None,
    *,
    page_size: int | None = None,
    max_scan: int | None = None,
    save: bool = True,
    snapshot_name: str | None = None,
    session: requests.Session | None = None,
) -> dict[str, Any]:
    """Resolve the top ``limit`` trending conversational models with live providers.

    The Hub API can return models whose only provider is staging or whose
    provider mapping is empty, so we scan more than ``limit`` rows and keep
    only candidates that have at least one provider marked ``live``.
    """
    if limit is None:
        limit = DEFAULT_MODEL_LIMIT
    if page_size is None:
        page_size = DEFAULT_PAGE_SIZE
    if max_scan is None:
        max_scan = DEFAULT_MAX_SCAN
    if limit < 1:
        raise ValueError("limit must be >= 1")

    client = session or requests.Session()
    headers = _auth_headers()
    models: list[dict[str, Any]] = []
    skipped: list[dict[str, Any]] = []
    scanned = 0
    pages = 0
    offset = 0

    while len(models) < limit and scanned < max_scan:
        params = {
            "inference_provider": "all",
            "filter": "conversational",
            "sort": "trendingScore",
            "direction": "-1",
            "limit": page_size,
            "offset": offset,
            "expand[]": "inferenceProviderMapping",
        }
        response = client.get(HUB_MODELS_API, params=params, headers=headers, timeout=30)
        response.raise_for_status()
        page = response.json()
        if not isinstance(page, list) or not page:
            break

        pages += 1
        for raw in page:
            if not isinstance(raw, dict):
                continue
            scanned += 1
            model_id = raw.get("id")
            if not model_id:
                continue

            live, statuses = _parse_provider_mapping(raw.get("inferenceProviderMapping"))
            if not live:
                skipped.append({"model": model_id, "provider_statuses": statuses})
                continue

            models.append(
                {
                    "rank": len(models) + 1,
                    "model": model_id,
                    "trending_score": raw.get("trendingScore"),
                    "downloads": raw.get("downloads"),
                    "pipeline_tag": raw.get("pipeline_tag"),
                    "live_providers": live,
                    "provider_statuses": statuses,
                }
            )
            if len(models) >= limit:
                break

        if len(page) < page_size:
            break
        offset += page_size

    if len(models) < limit:
        raise RuntimeError(
            f"Only found {len(models)} models with a live provider after scanning "
            f"{scanned} rows; requested {limit}. Increase SMOKE_MAX_SCAN or lower "
            "SMOKE_MODEL_LIMIT."
        )

    snapshot = {
        "fetched_at": utc_now(),
        "source": HUB_MODELS_API,
        "query": {
            "inference_provider": "all",
            "filter": "conversational",
            "sort": "trendingScore",
            "direction": "-1",
            "page_size": page_size,
            "max_scan": max_scan,
        },
        "requested_limit": limit,
        "returned": len(models),
        "scanned": scanned,
        "pages": pages,
        "skipped_without_live_provider": skipped,
        "models": models,
    }

    if save:
        name = snapshot_name or f"snapshot-{uuid.uuid4().hex[:8]}"
        _write_json(MODELS_DIR / f"{name}.json", snapshot)
        _write_json(MODELS_DIR / "latest.json", snapshot)

    return snapshot


def get_token_status(token: str | None = None, *, session: requests.Session | None = None) -> dict[str, Any]:
    """Validate HF_TOKEN and report the authenticated account and plan."""
    resolved = token or os.getenv("HF_TOKEN")
    if not resolved:
        return {"ok": False, "name": None, "plan": None, "error": "HF_TOKEN is not set."}

    client = session or requests.Session()
    try:
        response = client.get(
            WHOAMI_API,
            headers={"Authorization": f"Bearer {resolved}"},
            timeout=15,
        )
    except requests.RequestException as exc:
        return {"ok": False, "name": None, "plan": None, "error": f"Token check failed: {exc}"}

    if response.status_code != 200:
        detail = response.text.strip().replace("\n", " ")[:200]
        return {
            "ok": False,
            "name": None,
            "plan": None,
            "error": f"Token check returned HTTP {response.status_code}: {detail}",
        }

    data = response.json()
    return {
        "ok": True,
        "name": data.get("name"),
        "plan": data.get("plan"),
        "error": None,
    }


def load_latest_model_snapshot() -> dict[str, Any] | None:
    """Load the most recently saved discovery snapshot, if one exists."""
    path = MODELS_DIR / "latest.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None


_thread_state = threading.local()


def _client_for(token: str, timeout: float) -> InferenceClient:
    """Return one InferenceClient per worker thread.

    ``InferenceClient`` wraps a requests session, so sharing a single instance
    across worker threads is avoidable risk.
    """
    key = (token, timeout)
    if getattr(_thread_state, "key", None) != key:
        _thread_state.key = key
        _thread_state.client = InferenceClient(token=token, timeout=timeout)
    return _thread_state.client


def _format_exception(exc: BaseException, token: str) -> str:
    if isinstance(exc, HfHubHTTPError):
        response = getattr(exc, "response", None)
        status = getattr(response, "status_code", None)
        detail = getattr(response, "text", "") or str(exc)
        message = f"HTTP {status}: {detail}" if status else str(exc)
    elif isinstance(exc, InferenceTimeoutError):
        message = "Request timed out."
    else:
        message = f"{type(exc).__name__}: {exc}"

    message = " ".join(message.split())
    if token:
        message = message.replace(token, "hf_***")
    return message[:500]


def _http_status(exc: BaseException) -> int | None:
    response = getattr(exc, "response", None)
    return getattr(response, "status_code", None)


def _run_attempt(
    model_id: str,
    policy: str,
    token: str,
    timeout: float,
) -> dict[str, Any]:
    """Run a single chat-completion attempt and normalize its outcome."""
    started = time.perf_counter()
    try:
        response = _client_for(token, timeout).chat_completion(
            messages=[{"role": "user", "content": PROMPT}],
            model=f"{model_id}:{policy}",
            max_tokens=MAX_TOKENS,
            temperature=0.0,
            stream=False,
        )
    except Exception as exc:  # noqa: BLE001 - the point is to classify any failure
        return {
            "policy": policy,
            "ok": False,
            "latency_ms": round((time.perf_counter() - started) * 1000),
            "http_status": _http_status(exc),
            "error": _format_exception(exc, token),
            "usage": None,
        }

    latency_ms = round((time.perf_counter() - started) * 1000)
    choices = getattr(response, "choices", None) or []
    usage = getattr(response, "usage", None)
    usage_payload = (
        {
            "prompt_tokens": getattr(usage, "prompt_tokens", None),
            "completion_tokens": getattr(usage, "completion_tokens", None),
            "total_tokens": getattr(usage, "total_tokens", None),
        }
        if usage is not None
        else None
    )

    if not choices:
        return {
            "policy": policy,
            "ok": False,
            "latency_ms": latency_ms,
            "http_status": 200,
            "error": "Response contained no choices.",
            "usage": usage_payload,
        }

    return {
        "policy": policy,
        "ok": True,
        "latency_ms": latency_ms,
        "http_status": 200,
        "error": None,
        "usage": usage_payload,
    }


def test_model(
    candidate: dict[str, Any],
    *,
    token: str,
    timeout: float | None = None,
) -> dict[str, Any]:
    """Test one model with cheapest routing, then fastest on failure."""
    if timeout is None:
        timeout = DEFAULT_TIMEOUT
    attempts: list[dict[str, Any]] = []
    for policy in ATTEMPT_POLICIES:
        try:
            attempt = _run_attempt(candidate["model"], policy, token, timeout)
        except Exception as exc:  # noqa: BLE001 - defensive boundary for worker threads
            attempt = {
                "policy": policy,
                "ok": False,
                "latency_ms": 0,
                "http_status": None,
                "error": _format_exception(exc, token),
                "usage": None,
            }
        attempts.append(attempt)
        if attempt["ok"]:
            return {
                "rank": candidate["rank"],
                "model": candidate["model"],
                "ok": True,
                "attempt": policy,
                "fallback": policy == "fastest",
                "latency_ms": sum(item["latency_ms"] for item in attempts),
                "http_status": attempt["http_status"],
                "error": None,
                "attempts": attempts,
            }

    last = attempts[-1]
    return {
        "rank": candidate["rank"],
        "model": candidate["model"],
        "ok": False,
        "attempt": last["policy"],
        "fallback": True,
        "latency_ms": sum(item["latency_ms"] for item in attempts),
        "http_status": last["http_status"],
        "error": last["error"],
        "attempts": attempts,
    }


def run_smoke(
    models: list[dict[str, Any]],
    *,
    token: str | None = None,
    concurrency: int | None = None,
    timeout: float | None = None,
    run_id: str | None = None,
    progress: Callable[[int, int, dict[str, Any]], None] | None = None,
) -> dict[str, Any]:
    """Run the smoke test and persist a JSON result set."""
    if concurrency is None:
        concurrency = DEFAULT_CONCURRENCY
    if timeout is None:
        timeout = DEFAULT_TIMEOUT
    resolved_token = token or os.getenv("HF_TOKEN")
    if not resolved_token:
        raise RuntimeError("HF_TOKEN is not set; refusing to make inference calls.")
    if not models:
        raise ValueError("No models supplied for the smoke test.")

    resolved_run_id = run_id or new_run_id()
    if not RUN_ID_RE.match(resolved_run_id):
        raise ValueError("Invalid run_id.")

    started_at = utc_now()
    started_clock = time.perf_counter()
    results: list[dict[str, Any]] = []
    total = len(models)

    with ThreadPoolExecutor(max_workers=max(1, concurrency)) as pool:
        futures = {
            pool.submit(test_model, candidate, token=resolved_token, timeout=timeout): candidate
            for candidate in models
        }
        for done, future in enumerate(as_completed(futures), start=1):
            result = future.result()
            results.append(result)
            if progress is not None:
                progress(done, total, result)

    results.sort(key=lambda item: item["rank"])
    finished_at = utc_now()
    payload = {
        "run_id": resolved_run_id,
        "started_at": started_at,
        "finished_at": finished_at,
        "duration_seconds": round(time.perf_counter() - started_clock, 2),
        "model_limit": total,
        "passed": sum(1 for item in results if item["ok"]),
        "failed": sum(1 for item in results if not item["ok"]),
        "fallbacks": sum(1 for item in results if item["ok"] and item["fallback"]),
        "concurrency": concurrency,
        "max_tokens": MAX_TOKENS,
        "request_timeout_seconds": timeout,
        "results": results,
    }

    _write_json(RUNS_DIR / f"{resolved_run_id}.json", payload)
    _write_json(RESULTS_ROOT / "latest.json", payload)
    return payload


def list_runs(limit: int = 25) -> list[dict[str, Any]]:
    """Return recent run metadata, newest first."""
    ensure_storage()
    paths = sorted(RUNS_DIR.glob("*.json"), key=lambda path: path.stat().st_mtime, reverse=True)
    runs: list[dict[str, Any]] = []
    for path in paths[:limit]:
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        runs.append(
            {
                "run_id": payload.get("run_id", path.stem),
                "started_at": payload.get("started_at"),
                "passed": payload.get("passed"),
                "failed": payload.get("failed"),
                "model_limit": payload.get("model_limit"),
                "duration_seconds": payload.get("duration_seconds"),
            }
        )
    return runs


def load_run(run_id: str) -> dict[str, Any] | None:
    """Load one persisted run, rejecting path-traversal attempts."""
    if not RUN_ID_RE.match(run_id):
        raise ValueError("Invalid run_id.")
    path = RUNS_DIR / f"{run_id}.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None


def load_latest_run() -> dict[str, Any] | None:
    """Load the most recent completed run, if one exists."""
    path = RESULTS_ROOT / "latest.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None