from __future__ import annotations

import json
import os
import tempfile
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path

from health_client import AccountSnapshot, HealthStatus
from limit_warning import WINDOW_5H, WINDOW_7D
from runtime_paths import runtime_dir
from spend_cap import capped_windows, describe_caps

_VALID_CAP_WINDOWS = frozenset({WINDOW_5H, WINDOW_7D})


@dataclass(frozen=True)
class RoutingRules:
    projects: dict[str, "ProjectRoute"]
    default: str
    quota_exhausted_threshold_pct: int
    fallback_chain: tuple[str, ...] = ()
    fallback_trigger: str = "broken_or_quota_exhausted"
    missing_health_is_available: bool = True
    account_caps: Mapping[str, Mapping[str, int]] = field(default_factory=dict)


@dataclass(frozen=True)
class ProjectRoute:
    account: str
    fallback: tuple[str, ...] = ()


@dataclass(frozen=True)
class ResolvedRoute:
    slug: str
    chain: tuple[str, ...]
    source: str
    fallback_used: bool
    fallback_from: str | None


class NoHealthyAccountError(RuntimeError):
    def __init__(self, message: str, *, chain: tuple[str, ...] = ()) -> None:
        super().__init__(message)
        self.chain = chain


class AllAccountsCappedError(Exception):
    def __init__(self, resume_at: str | None) -> None:
        self.resume_at = resume_at


class RoutingResolver:
    def __init__(
        self,
        rules: RoutingRules,
        health: dict[str, AccountSnapshot],
        known_slugs: set[str] | None = None,
        locked_slugs: set[str] | frozenset[str] | None = None,
    ) -> None:
        self.rules = rules
        self.health = health
        self.known_slugs = None if known_slugs is None else frozenset(known_slugs)
        self.locked_slugs = frozenset(locked_slugs or ())

    def resolve(self, project_name: str | None, dynamic_slug: str | None = None) -> ResolvedRoute:
        project_route = self.rules.projects.get(project_name) if project_name is not None else None
        if isinstance(project_route, str):
            project_route = ProjectRoute(project_route)
        source = "projects" if project_route is not None else "default"
        raw_chain = ([project_route.account, *project_route.fallback]
                     if project_route is not None else [self.rules.default, *self.rules.fallback_chain])
        chain_parts: list[str] = []
        for slug in raw_chain:
            if slug == "dynamic":
                if not dynamic_slug:
                    raise NoHealthyAccountError("dynamic route has no selected account")
                chain_parts.append(dynamic_slug)
            else:
                chain_parts.append(slug)
        chain = tuple(chain_parts)
        for index, slug in enumerate(chain):
            if self._is_available(slug):
                return ResolvedRoute(slug, chain, source, index > 0, chain[0] if index > 0 else None)
        raise NoHealthyAccountError(self._format_error(list(chain)), chain=chain)

    def account_caps(self, slug: str) -> Mapping[str, int]:
        return self.rules.account_caps.get(slug, {})

    def _is_cap_excluded(self, slug: str) -> bool:
        caps = self.account_caps(slug)
        if not caps:
            return False
        snapshot = self.health.get(slug)
        if snapshot is None:
            return True
        if snapshot.status == HealthStatus.BROKEN:
            return False
        return bool(capped_windows(snapshot, caps))

    def all_cap_excluded(self, slugs: tuple[str, ...]) -> bool:
        if not slugs:
            return False
        return all(self._is_cap_excluded(slug) for slug in slugs)

    def all_locked(self, slugs: tuple[str, ...]) -> bool:
        return bool(slugs) and all(slug in self.locked_slugs for slug in slugs)

    def unlocked(self, slugs: tuple[str, ...]) -> tuple[str, ...]:
        return tuple(slug for slug in slugs if slug not in self.locked_slugs)

    def earliest_cap_resume_at(self, slugs: tuple[str, ...]) -> str | None:
        epochs: list[float] = []
        for slug in slugs:
            caps = self.account_caps(slug)
            snapshot = self.health.get(slug)
            if not caps or snapshot is None:
                continue
            for window in capped_windows(snapshot, caps):
                reset_at = (
                    snapshot.primary_reset_at
                    if window == WINDOW_5H
                    else snapshot.secondary_reset_at
                )
                if reset_at is not None:
                    epochs.append(reset_at)
        if not epochs:
            return None
        dt = datetime.fromtimestamp(int(min(epochs)), tz=UTC)
        return dt.strftime("%Y-%m-%dT%H:%M:%SZ")

    def _is_available(self, slug: str) -> bool:
        if slug in self.locked_slugs:
            return False
        if self.known_slugs is not None and slug not in self.known_slugs:
            return False
        snapshot = self.health.get(slug)
        if snapshot is None:
            if self.account_caps(slug):
                return False
            return self.known_slugs is not None and self.rules.missing_health_is_available
        if snapshot.status == HealthStatus.BROKEN:
            return False
        if capped_windows(snapshot, self.account_caps(slug)):
            return False
        if self.rules.fallback_trigger != "broken_or_quota_exhausted":
            return True
        return not self._quota_exhausted(snapshot)

    def _quota_exhausted(self, snapshot: AccountSnapshot) -> bool:
        threshold = self.rules.quota_exhausted_threshold_pct
        primary = snapshot.primary_used_pct
        secondary = snapshot.secondary_used_pct
        return (primary is not None and primary >= threshold) or (
            secondary is not None and secondary >= threshold
        )

    def _format_error(self, tried: list[str]) -> str:
        parts = [f"{slug}={self._describe_status(slug)}" for slug in tried]
        return "no healthy account available: " + ", ".join(parts)

    def _describe_status(self, slug: str) -> str:
        if slug in self.locked_slugs:
            return "locked"
        if self.known_slugs is not None and slug not in self.known_slugs:
            return "unknown-account"
        snapshot = self.health.get(slug)
        if snapshot is None:
            return "missing-health"
        if snapshot.status == HealthStatus.BROKEN:
            return "broken"
        cap_status = describe_caps(snapshot, self.account_caps(slug))
        if cap_status:
            return cap_status
        if (
            self.rules.fallback_trigger == "broken_or_quota_exhausted"
            and self._quota_exhausted(snapshot)
        ):
            primary = "?" if snapshot.primary_used_pct is None else str(snapshot.primary_used_pct)
            secondary = (
                "?" if snapshot.secondary_used_pct is None else str(snapshot.secondary_used_pct)
            )
            return f"quota-exhausted(primary={primary},secondary={secondary})"
        return snapshot.status.value


def _normalize_account_caps(
    account_caps: object,
    *,
    known_slugs: set[str] | None,
) -> dict[str, dict[str, int]]:
    if account_caps is None:
        return {}
    if not isinstance(account_caps, dict):
        raise ValueError("routing_rules.json account_caps must be an object")
    normalized: dict[str, dict[str, int]] = {}
    for slug, windows in account_caps.items():
        if not isinstance(slug, str) or not slug:
            raise ValueError("routing_rules.json account_caps keys must be non-empty strings")
        if known_slugs is not None and slug not in known_slugs:
            raise ValueError(f"routing_rules.json account_caps slug {slug!r} is not in the registry")
        if not isinstance(windows, dict):
            raise ValueError("routing_rules.json account_caps entries must be objects")
        normalized_windows: dict[str, int] = {}
        for window, cap_pct in windows.items():
            if window not in _VALID_CAP_WINDOWS:
                raise ValueError("routing_rules.json account_caps window keys must be '5h' or '7d'")
            if not isinstance(cap_pct, int) or not 1 <= cap_pct <= 100:
                raise ValueError("routing_rules.json account_caps values must be ints from 1 to 100")
            normalized_windows[window] = cap_pct
        normalized[slug] = normalized_windows
    return normalized


def load_rules(
    path: Path | None = None,
    *,
    known_slugs: set[str] | None = None,
) -> RoutingRules:
    if path is None:
        path = runtime_dir() / "routing_rules.json"
    payload = json.loads(path.read_text(encoding="utf-8"))
    projects = payload.get("projects")
    if not isinstance(projects, dict):
        raise ValueError("routing_rules.json projects must be an object")

    default = payload.get("default")
    fallback_chain = payload.get("fallback_chain", [])
    fallback_trigger = payload.get("fallback_trigger", "broken_or_quota_exhausted")
    missing_health_is_available = payload.get(
        "missing_health_is_available", payload.get("version") != "routing/v2"
    )
    threshold = payload.get("quota_exhausted_threshold_pct", 100)
    if not isinstance(default, str) or not default:
        raise ValueError("routing_rules.json default must be a non-empty string")
    if not isinstance(threshold, int):
        raise ValueError("routing_rules.json quota_exhausted_threshold_pct must be an int")
    if not isinstance(fallback_chain, list) or any(
        not isinstance(slug, str) or not slug for slug in fallback_chain
    ):
        raise ValueError("routing_rules.json fallback_chain must be a list of non-empty strings")
    if not isinstance(fallback_trigger, str) or not fallback_trigger:
        raise ValueError("routing_rules.json fallback_trigger must be a non-empty string")
    if not isinstance(missing_health_is_available, bool):
        raise ValueError("routing_rules.json missing_health_is_available must be a bool")

    normalized_projects: dict[str, ProjectRoute] = {}
    for key, value in projects.items():
        if isinstance(value, str):
            value = {"account": value}
        if not isinstance(key, str) or not key or not isinstance(value, dict):
            raise ValueError("routing_rules.json projects entries must be route objects")
        account, fallback = value.get("account"), value.get("fallback", [])
        if not isinstance(account, str) or not account or not isinstance(fallback, list) or any(not isinstance(x, str) or not x for x in fallback):
            raise ValueError("routing_rules.json project route is invalid")
        normalized_projects[key] = ProjectRoute(account, tuple(fallback))

    account_caps = _normalize_account_caps(
        payload.get("account_caps", {}),
        known_slugs=known_slugs,
    )

    if payload.get("version") != "routing/v2":
        migrated = {
            "version": "routing/v2",
            "projects": {
                key: {
                    "account": route.account,
                    **({"fallback": list(route.fallback)} if route.fallback else {}),
                }
                for key, route in normalized_projects.items()
            },
            "default": default,
            "fallback_chain": fallback_chain,
            "fallback_trigger": fallback_trigger,
            "missing_health_is_available": missing_health_is_available,
            "quota_exhausted_threshold_pct": threshold,
            "account_caps": account_caps,
        }
        path.write_text(json.dumps(migrated, indent=2) + "\n", encoding="utf-8")

    return RoutingRules(
        projects=normalized_projects,
        default=default,
        fallback_chain=tuple(fallback_chain),
        fallback_trigger=fallback_trigger,
        missing_health_is_available=missing_health_is_available,
        quota_exhausted_threshold_pct=threshold,
        account_caps=account_caps,
    )


def _atomic_write_rules(path: Path, payload: dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        "w",
        encoding="utf-8",
        dir=path.parent,
        delete=False,
    ) as temp_file:
        json.dump(payload, temp_file, indent=2, sort_keys=True)
        temp_file.write("\n")
        temp_file.flush()
        os.fsync(temp_file.fileno())
        temp_path = Path(temp_file.name)
    temp_path.replace(path)


def update_account_cap(
    path: Path,
    slug: str,
    window: str,
    cap_pct: int | None,
    *,
    known_slugs: set[str] | None = None,
) -> None:
    try:
        raw = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        raise ValueError(f"routing rules unreadable: {path}") from exc
    if not isinstance(raw, dict):
        raise ValueError(f"routing rules unreadable: {path}")

    account_caps = raw.get("account_caps", {})
    if account_caps is None:
        account_caps = {}
    if not isinstance(account_caps, dict):
        raise ValueError("routing_rules.json account_caps must be an object")

    slug_caps = account_caps.get(slug, {})
    if slug_caps is None:
        slug_caps = {}
    if not isinstance(slug_caps, dict):
        raise ValueError("routing_rules.json account_caps entries must be objects")

    updated_slug_caps = dict(slug_caps)
    if cap_pct is None:
        updated_slug_caps.pop(window, None)
    else:
        updated_slug_caps[window] = cap_pct

    updated_account_caps = dict(account_caps)
    if known_slugs is not None:
        updated_account_caps = {
            entry_slug: windows
            for entry_slug, windows in updated_account_caps.items()
            if entry_slug in known_slugs
        }
    if updated_slug_caps:
        updated_account_caps[slug] = updated_slug_caps
    else:
        updated_account_caps.pop(slug, None)

    _normalize_account_caps(updated_account_caps, known_slugs=known_slugs)

    payload = dict(raw)
    payload["account_caps"] = updated_account_caps
    _atomic_write_rules(path, payload)