from __future__ import annotations

import json
import os
import shutil
import subprocess
import tempfile
import time
from dataclasses import dataclass, replace
from datetime import datetime
from pathlib import Path

from account_registry import Account
from claude_account_activity import account_has_live_session
from claude_credentials import (
    credentials_fingerprint,
    pin_account_uuid,
    pinned_account_uuid,
    resolve_credentials,
)
from claude_identity import ClaudeTokenIdentity
from claude_oauth import (
    ClaudeOAuthAuthError,
    ClaudeOAuthError,
    ClaudeOAuthLoggedOutError,
    ClaudeOAuthSession,
    ClaudeOAuthTransportError,
    UsageReading,
)
from health_client import (
    AccountSnapshot,
    ExtraUsageSummary,
    HealthStatus,
    NamedLimit,
    SpendSummary,
)

CLAUDE_CACHE_FRESH_SECONDS = 10 * 60
CLAUDE_DEFAULT_RETRY_SECONDS = 5 * 60
CLAUDE_MIN_RETRY_SECONDS = 60
CLAUDE_LOGGED_OUT_RETRY_SECONDS = 30 * 60
CLAUDE_WINDOW_SECONDS = {"primary": 5 * 60 * 60, "secondary": 7 * 24 * 60 * 60}
# Every successful rotation drops a restorable copy of the fresh grant here; for
# accounts whose grant is the only access, this copy is the recovery path.
ROTATION_BACKUP_DIR = Path.home() / ".local/state/overdeck/token-rescue"


@dataclass(frozen=True)
class _Backoff:
    retry_not_before: float
    logged_out: bool


class ClaudeHealthClient:
    def fetch(self, account: Account, timeout_secs: float = 10.0) -> AccountSnapshot:
        checked_at = time.time()
        identity = ClaudeTokenIdentity(timeout_secs=timeout_secs)
        own_path = account.account_home / ".credentials.json"
        if pinned_account_uuid(account.account_home) is None:
            owner_uuid = identity.account_uuid_for(own_path)
            if owner_uuid is not None:
                pin_account_uuid(account.account_home, owner_uuid)
        credentials = resolve_credentials(
            account.account_home,
            identify=identity.account_uuid_for,
        )
        auth_status = self._auth_status(account.account_home, credentials.path, timeout_secs)
        if auth_status == HealthStatus.BROKEN:
            return AccountSnapshot(
                status=HealthStatus.BROKEN,
                primary_used_pct=None,
                secondary_used_pct=None,
                checked_at=checked_at,
                detail="authentication required",
            )
        fingerprint = credentials_fingerprint(credentials.path)
        backoff = self._read_backoff(account.account_home, fingerprint)
        if checked_at < backoff.retry_not_before:
            if backoff.logged_out:
                return AccountSnapshot(
                    status=HealthStatus.BROKEN,
                    primary_used_pct=None,
                    secondary_used_pct=None,
                    checked_at=checked_at,
                    detail="authentication required",
                )
            return self._cached_snapshot(
                account.account_home,
                checked_at,
                status_override=(
                    HealthStatus.UNKNOWN if auth_status == HealthStatus.UNKNOWN else None
                ),
                unavailable_detail=(
                    None if auth_status == HealthStatus.UNKNOWN else "limits unavailable"
                ),
            )

        # Rotation is only safe when this poller is provably the sole grant holder:
        # the account's own file (never the shared live one), no CLI session running
        # against that config dir.
        allow_rotation = (
            credentials.path == account.account_home / ".credentials.json"
            and not account_has_live_session(account.account_home)
        )
        try:
            reading = ClaudeOAuthSession(
                credentials.path,
                timeout_secs=timeout_secs,
                allow_rotation=allow_rotation,
                rotation_backup_path=ROTATION_BACKUP_DIR
                / f"{account.ref.slug}.credentials.json.auto",
            ).get_usage_reading()
        except ClaudeOAuthLoggedOutError:
            # `claude auth status` only parses the local credentials file, so a server-side
            # token rejection outranks it.
            self._write_retry_not_before(
                account.account_home,
                checked_at + CLAUDE_LOGGED_OUT_RETRY_SECONDS,
                fingerprint=fingerprint,
                logged_out=True,
            )
            return AccountSnapshot(
                status=HealthStatus.BROKEN,
                primary_used_pct=None,
                secondary_used_pct=None,
                checked_at=checked_at,
                detail="authentication required",
            )
        except ClaudeOAuthAuthError:
            if auth_status != HealthStatus.OK:
                return AccountSnapshot(
                    status=HealthStatus.BROKEN,
                    primary_used_pct=None,
                    secondary_used_pct=None,
                    checked_at=checked_at,
                    detail="authentication required",
                )
            self._write_retry_not_before(
                account.account_home,
                checked_at + CLAUDE_DEFAULT_RETRY_SECONDS,
                fingerprint=fingerprint,
            )
            return self._cached_snapshot(account.account_home, checked_at)
        except ClaudeOAuthError as exc:
            retry_seconds = (
                max(exc.retry_after_seconds, CLAUDE_MIN_RETRY_SECONDS)
                if isinstance(exc, ClaudeOAuthTransportError)
                and exc.retry_after_seconds is not None
                else CLAUDE_DEFAULT_RETRY_SECONDS
            )
            self._write_retry_not_before(
                account.account_home,
                checked_at + retry_seconds,
                fingerprint=fingerprint,
                detail=f"{type(exc).__name__}: {exc}",
            )
            return self._cached_snapshot(
                account.account_home,
                checked_at,
                status_override=(
                    HealthStatus.UNKNOWN if auth_status == HealthStatus.UNKNOWN else None
                ),
                unavailable_detail=(
                    None if auth_status == HealthStatus.UNKNOWN else "limits unavailable"
                ),
            )

        # The shared live credentials file can be swapped to another account mid-fetch,
        # so a reading only counts when its token provably belongs to this account.
        if not self._reading_owner_matches(account.account_home, reading, identity):
            return self._cached_snapshot(
                account.account_home,
                checked_at,
                unavailable_detail="usage attribution unverified",
            )
        self._clear_retry_not_before(account.account_home)
        self._write_rate_limits_payload(account.account_home, reading.payload)
        return self._snapshot_from_payload(reading.payload, HealthStatus.OK, checked_at)

    @staticmethod
    def _reading_owner_matches(
        account_home: Path,
        reading: UsageReading,
        identity: ClaudeTokenIdentity,
    ) -> bool:
        expected_uuid = pinned_account_uuid(account_home)
        if expected_uuid is None:
            return False
        owner = identity.identify_token(reading.access_token)
        return owner is not None and owner.account_uuid == expected_uuid

    @classmethod
    def _snapshot_from_payload(
        cls,
        payload: dict[str, object],
        status: HealthStatus,
        checked_at: float,
        now: float | None = None,
    ) -> AccountSnapshot:
        now = checked_at if now is None else now
        named_limits = cls._current_limits(cls._named_limits(payload), now)
        primary = cls._limit_for_group(named_limits, "primary")
        secondary = cls._limit_for_group(named_limits, "secondary")
        return AccountSnapshot(
            status=status,
            primary_used_pct=primary.percent if primary is not None else None,
            secondary_used_pct=secondary.percent if secondary is not None else None,
            primary_reset_at=cls._reset_or_default(
                primary, checked_at, CLAUDE_WINDOW_SECONDS["primary"]
            ),
            secondary_reset_at=cls._reset_or_default(
                secondary, checked_at, CLAUDE_WINDOW_SECONDS["secondary"]
            ),
            checked_at=checked_at,
            named_limits=tuple(named_limits),
            extra_usage=cls._parse_extra_usage(
                cls._first_value(
                    payload,
                    "extra_usage",
                    "extraUsage",
                    "extra_usage_summary",
                    "extraUsageSummary",
                )
            ),
            spend=cls._parse_spend(
                cls._first_value(
                    payload,
                    "spend",
                    "spendSummary",
                    "spend_summary",
                    "extra_spend",
                    "extraSpend",
                )
            ),
        )

    @classmethod
    def _cached_snapshot(
        cls,
        account_home: Path,
        now: float,
        *,
        status_override: HealthStatus | None = None,
        unavailable_detail: str | None = "limits unavailable",
    ) -> AccountSnapshot:
        path = account_home / "rate-limits-cache.json"
        payload = cls._read_rate_limits_payload(account_home)
        if payload is None:
            return AccountSnapshot(
                status=HealthStatus.UNKNOWN,
                primary_used_pct=None,
                secondary_used_pct=None,
                checked_at=now,
                detail=unavailable_detail,
            )
        try:
            checked_at = min(path.stat().st_mtime, now)
        except OSError:
            checked_at = now
        age = max(0.0, now - checked_at)
        status = status_override or (
            HealthStatus.OK if age <= CLAUDE_CACHE_FRESH_SECONDS else HealthStatus.UNKNOWN
        )
        return cls._snapshot_from_payload(payload, status, checked_at, now=now)

    @classmethod
    def _current_limits(cls, limits: list[NamedLimit], now: float) -> list[NamedLimit]:
        return [cls._current_limit(limit, now) for limit in limits]

    @staticmethod
    def _current_limit(limit: NamedLimit, now: float) -> NamedLimit:
        """A window's counter zeroes at ``resets_at``, so a reading carried past that
        boundary describes a window that no longer exists."""
        if limit.resets_at is None or limit.resets_at > now:
            return limit
        period = CLAUDE_WINDOW_SECONDS.get(limit.group or "")
        if period is None:
            return replace(limit, percent=0, resets_at=None)
        elapsed_windows = int((now - limit.resets_at) // period) + 1
        return replace(limit, percent=0, resets_at=limit.resets_at + elapsed_windows * period)

    @staticmethod
    def _reset_or_default(
        limit: NamedLimit | None,
        checked_at: float,
        default_seconds: int,
    ) -> float | None:
        if limit is None or limit.percent is None:
            return None
        if limit.resets_at is not None:
            return limit.resets_at
        return checked_at + default_seconds

    @staticmethod
    def _auth_status(
        account_home: Path,
        credentials_path: Path,
        timeout_secs: float,
    ) -> HealthStatus:
        try:
            with tempfile.TemporaryDirectory(prefix="claude-health-") as temp_dir:
                config_dir = Path(temp_dir)
                if credentials_path.exists():
                    shutil.copy2(credentials_path, config_dir / ".credentials.json")
                for name in ("claude.json", ".claude.json"):
                    source = account_home / name
                    if source.exists():
                        shutil.copy2(source, config_dir / name)
                result = subprocess.run(
                    ["claude", "auth", "status", "--json"],
                    check=False,
                    capture_output=True,
                    text=True,
                    env={**os.environ, "CLAUDE_CONFIG_DIR": str(config_dir)},
                    timeout=timeout_secs,
                )
        except (OSError, subprocess.TimeoutExpired):
            return HealthStatus.UNKNOWN

        stdout = "" if result.stdout is None else result.stdout.strip()
        stderr = "" if result.stderr is None else result.stderr.strip()
        try:
            payload = json.loads(stdout)
        except json.JSONDecodeError:
            if result.returncode != 0 and ClaudeHealthClient._is_auth_rejection(stdout, stderr):
                return HealthStatus.BROKEN
            return HealthStatus.UNKNOWN

        if not isinstance(payload, dict):
            return HealthStatus.UNKNOWN
        logged_in = payload.get("loggedIn")
        if logged_in is True:
            return HealthStatus.OK
        if logged_in is False:
            return HealthStatus.BROKEN
        if result.returncode != 0 and ClaudeHealthClient._is_auth_rejection(stdout, stderr):
            return HealthStatus.BROKEN
        return HealthStatus.UNKNOWN

    @staticmethod
    def _is_auth_rejection(stdout: str, stderr: str) -> bool:
        haystack = f"{stdout}\n{stderr}".lower()
        return any(
            token in haystack
            for token in (
                "authentication required",
                "login required",
                "not logged in",
                "logged out",
                "unauthorized",
                "forbidden",
                "invalid credentials",
                "missing credentials",
            )
        )

    @staticmethod
    def _retry_path(account_home: Path) -> Path:
        return account_home / "usage-backoff.json"

    @classmethod
    def _read_backoff(cls, account_home: Path, fingerprint: str | None) -> _Backoff:
        try:
            payload = json.loads(cls._retry_path(account_home).read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            return _Backoff(0.0, False)
        if not isinstance(payload, dict):
            return _Backoff(0.0, False)
        if payload.get("credentials_fingerprint") != fingerprint:
            return _Backoff(0.0, False)
        value = payload.get("retry_not_before")
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            return _Backoff(0.0, False)
        return _Backoff(float(value), payload.get("logged_out") is True)

    @classmethod
    def _write_retry_not_before(
        cls,
        account_home: Path,
        retry_not_before: float,
        fingerprint: str | None = None,
        logged_out: bool = False,
        detail: str | None = None,
    ) -> None:
        try:
            cls._retry_path(account_home).write_text(
                json.dumps(
                    {
                        "retry_not_before": retry_not_before,
                        "credentials_fingerprint": fingerprint,
                        "logged_out": logged_out,
                        "detail": detail,
                    }
                ),
                encoding="utf-8",
            )
        except OSError:
            return

    @classmethod
    def _clear_retry_not_before(cls, account_home: Path) -> None:
        try:
            cls._retry_path(account_home).unlink(missing_ok=True)
        except OSError:
            return

    @staticmethod
    def _read_rate_limits_payload(account_home: Path) -> dict[str, object] | None:
        path = account_home / "rate-limits-cache.json"
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            return None
        if not isinstance(payload, dict):
            return None
        return payload

    @staticmethod
    def _write_rate_limits_payload(account_home: Path, payload: dict[str, object]) -> None:
        path = account_home / "rate-limits-cache.json"
        tmp_name: str | None = None
        try:
            fd, tmp_name = tempfile.mkstemp(dir=account_home, prefix=".rate-limits-cache.json.")
            with os.fdopen(fd, "w", encoding="utf-8") as handle:
                handle.write(json.dumps(payload, sort_keys=True))
            os.replace(tmp_name, path)
        except OSError:
            if tmp_name is not None:
                try:
                    os.unlink(tmp_name)
                except OSError:
                    pass

    @staticmethod
    def _named_limits(payload: dict[str, object]) -> list[NamedLimit]:
        limits: list[NamedLimit] = []
        seen: set[str] = set()
        raw_limits = ClaudeHealthClient._first_value(payload, "named_limits", "namedLimits", "limits")
        if isinstance(raw_limits, list):
            for raw_limit in raw_limits:
                if not isinstance(raw_limit, dict):
                    continue
                kind = raw_limit.get("kind", raw_limit.get("name"))
                limit = ClaudeHealthClient._parse_named_limit(
                    kind if isinstance(kind, str) else "",
                    raw_limit,
                )
                if limit is None or not limit.active or limit.kind in seen:
                    continue
                seen.add(limit.kind)
                limits.append(limit)
        for kind, value in payload.items():
            if kind in {"named_limits", "namedLimits", "limits", "extra_usage", "extraUsage", "extra_usage_summary", "extraUsageSummary", "spend", "spendSummary", "spend_summary", "extra_spend", "extraSpend"}:
                continue
            limit = ClaudeHealthClient._parse_named_limit(kind, value)
            if limit is None or not limit.active or limit.kind in seen:
                continue
            seen.add(limit.kind)
            limits.append(limit)
        return limits

    @staticmethod
    def _parse_named_limit(kind: str, value: object) -> NamedLimit | None:
        if not isinstance(kind, str) or not isinstance(value, dict):
            return None
        percent = ClaudeHealthClient._percent(
            ClaudeHealthClient._first_value(
                value,
                "percent",
                "pct",
                "usedPercent",
                "used_percentage",
                "utilization",
            )
        )
        resets_at = ClaudeHealthClient._reset_at(
            value.get("resets_at", value.get("resetsAt", value.get("resetAt")))
        )
        group = ClaudeHealthClient._group(kind, value.get("group"))
        active_value = ClaudeHealthClient._first_value(value, "active", "is_active")
        active = active_value if isinstance(active_value, bool) else True
        if percent is None and resets_at is None and group is None and active_value is None:
            return None
        return NamedLimit(kind=kind, group=group, percent=percent, resets_at=resets_at, active=active)

    @staticmethod
    def _group(kind: str, group: object) -> str | None:
        if kind == "five_hour":
            return "primary"
        if kind == "seven_day":
            return "secondary"
        return group if isinstance(group, str) and group else None

    @staticmethod
    def _limit_for_group(limits: list[NamedLimit], group: str) -> NamedLimit | None:
        for limit in limits:
            if limit.group == group:
                return limit
        return None

    @staticmethod
    def _percent(value: object) -> int | None:
        if isinstance(value, bool):
            return None
        if isinstance(value, (int, float)):
            return max(0, min(100, int(value)))
        return None

    @staticmethod
    def _reset_at(value: object) -> float | None:
        if isinstance(value, bool):
            return None
        if isinstance(value, (int, float)):
            return float(value)
        if isinstance(value, str):
            normalized = value.strip()
            if not normalized:
                return None
            if normalized.endswith("Z"):
                normalized = f"{normalized[:-1]}+00:00"
            try:
                return datetime.fromisoformat(normalized).timestamp()
            except ValueError:
                return None
        return None

    @staticmethod
    def _parse_extra_usage(value: object) -> ExtraUsageSummary | None:
        if not isinstance(value, dict):
            return None
        used = ClaudeHealthClient._float(value.get("used"))
        limit = ClaudeHealthClient._float(value.get("limit"))
        unit = value.get("unit")
        display = value.get("display")
        if used is None and limit is None and not isinstance(unit, str) and not isinstance(display, str):
            return None
        return ExtraUsageSummary(
            used=used,
            limit=limit,
            unit=unit if isinstance(unit, str) else None,
            display=display if isinstance(display, str) else None,
        )

    @staticmethod
    def _parse_spend(value: object) -> SpendSummary | None:
        if not isinstance(value, dict):
            return None
        amount = ClaudeHealthClient._float(value.get("amount"))
        limit = ClaudeHealthClient._float(value.get("limit"))
        currency = value.get("currency")
        display = value.get("display")
        if (
            amount is None
            and limit is None
            and not isinstance(currency, str)
            and not isinstance(display, str)
        ):
            return None
        return SpendSummary(
            amount=amount,
            limit=limit,
            currency=currency if isinstance(currency, str) else None,
            display=display if isinstance(display, str) else None,
        )

    @staticmethod
    def _float(value: object) -> float | None:
        if isinstance(value, bool):
            return None
        if isinstance(value, (int, float)):
            return float(value)
        return None

    @staticmethod
    def _first_value(payload: dict[str, object], *keys: str) -> object:
        for key in keys:
            if key in payload:
                return payload[key]
        return None