#!/usr/bin/env python3
"""Discover the current Claude subscription model catalog without inference.

Claude Code exposes the account-visible ``ModelInfo[]`` through its Agent SDK
``initialize`` control response. This module speaks that documented stream-JSON
control protocol directly instead of treating authentication status or a cached
model list as a live catalog. It intentionally retains only non-secret
ModelInfo metadata.
"""

from __future__ import annotations

import json
import os
import signal
import subprocess
import uuid
from dataclasses import dataclass
from typing import Any, Mapping, Sequence


DEFAULT_CLAUDE_BIN = "/opt/data/tools/bin/claude"
DEFAULT_TIMEOUT_SECONDS = 30.0
MAX_MODELS = 256
EFFORT_LEVELS = frozenset({"low", "medium", "high", "xhigh", "max"})


@dataclass(frozen=True)
class ClaudeModel:
    """Safe metadata for one model advertised by the native Claude CLI."""

    value: str
    resolved_model: str | None
    display_name: str
    description: str
    supports_effort: bool
    supported_effort_levels: tuple[str, ...]
    supports_adaptive_thinking: bool
    supports_fast_mode: bool
    supports_auto_mode: bool
    disabled: bool

    def as_dict(self) -> dict[str, Any]:
        """Serialize only ModelInfo metadata, preserving Claude's alias value."""
        return {
            "value": self.value,
            "resolvedModel": self.resolved_model,
            "displayName": self.display_name,
            "description": self.description,
            "supportsEffort": self.supports_effort,
            "supportedEffortLevels": list(self.supported_effort_levels),
            "supportsAdaptiveThinking": self.supports_adaptive_thinking,
            "supportsFastMode": self.supports_fast_mode,
            "supportsAutoMode": self.supports_auto_mode,
            "disabled": self.disabled,
        }


@dataclass(frozen=True)
class ClaudeModelDiscovery:
    """Result of one bounded, subscription-authenticated catalog probe."""

    models: tuple[ClaudeModel, ...]
    live: bool
    connected: bool
    state: str
    error_code: str | None = None


def _native_environment(base: Mapping[str, str] | None = None) -> dict[str, str]:
    """Keep mounted OAuth configuration while preventing API-key billing lanes."""
    environment = dict(os.environ if base is None else base)
    environment.pop("ANTHROPIC_API_KEY", None)
    environment.pop("CLAUDE_API_KEY", None)
    return environment


def _command(claude_bin: str) -> list[str]:
    """Build a no-inference CLI invocation with all tools and MCP disabled."""
    return [
        claude_bin,
        "--input-format",
        "stream-json",
        "--output-format",
        "stream-json",
        "--verbose",
        "--print",
        "--no-session-persistence",
        "--safe-mode",
        "--tools",
        "",
        "--strict-mcp-config",
        "--mcp-config",
        '{"mcpServers":{}}',
    ]


def _initialize_request() -> str:
    """Build the Claude Agent SDK initialize request for ModelInfo discovery."""
    return json.dumps(
        {
            "type": "control_request",
            "request_id": f"hermes-model-catalog-{uuid.uuid4()}",
            "request": {"subtype": "initialize", "hooks": {}},
        },
        separators=(",", ":"),
    )


def _stop_process(process: subprocess.Popen[str]) -> None:
    """End a timed-out CLI process and its process group without leaking output."""
    if process.poll() is not None:
        return
    try:
        os.killpg(process.pid, signal.SIGTERM)
    except (OSError, ProcessLookupError):
        process.terminate()
    try:
        process.communicate(timeout=2)
        return
    except subprocess.TimeoutExpired:
        pass
    try:
        os.killpg(process.pid, signal.SIGKILL)
    except (OSError, ProcessLookupError):
        process.kill()
    try:
        process.communicate(timeout=2)
    except subprocess.TimeoutExpired:
        pass


def _run_initialize(
    command: Sequence[str], environment: Mapping[str, str], timeout_seconds: float
) -> tuple[str | None, str | None]:
    """Run initialize once, returning stream output or a non-sensitive error code."""
    try:
        process = subprocess.Popen(
            list(command),
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
            env=dict(environment),
            start_new_session=True,
        )
    except OSError:
        return None, "spawn-failed"
    try:
        request = _initialize_request() + "\n"
        stdout, _ = process.communicate(request, timeout=timeout_seconds)
    except subprocess.TimeoutExpired:
        _stop_process(process)
        return None, "timeout"
    except OSError:
        _stop_process(process)
        return None, "transport-error"
    if process.returncode != 0:
        return None, "cli-error"
    return stdout, None


def _text(value: Any, *, maximum: int = 1024) -> str:
    """Return bounded string metadata and reject non-string protocol values."""
    return value[:maximum] if isinstance(value, str) else ""


def _model_info(value: Any) -> ClaudeModel | None:
    """Validate one ModelInfo record without inventing aliases or identifiers."""
    if not isinstance(value, dict):
        return None
    alias = _text(value.get("value"))
    if not alias:
        return None
    resolved_model = _text(value.get("resolvedModel")) or None
    efforts = value.get("supportedEffortLevels")
    supported_effort_levels = (
        tuple(
            level
            for level in efforts
            if isinstance(level, str) and level in EFFORT_LEVELS
        )
        if isinstance(efforts, list)
        else ()
    )
    return ClaudeModel(
        value=alias,
        resolved_model=resolved_model,
        display_name=_text(value.get("displayName")),
        description=_text(value.get("description")),
        supports_effort=value.get("supportsEffort") is True,
        supported_effort_levels=supported_effort_levels,
        supports_adaptive_thinking=value.get("supportsAdaptiveThinking") is True,
        supports_fast_mode=value.get("supportsFastMode") is True,
        supports_auto_mode=value.get("supportsAutoMode") is True,
        disabled=value.get("disabled") is True,
    )


def _parse_initialize_output(stdout: str) -> ClaudeModelDiscovery:
    """Extract all ModelInfo pages present in initialize control responses.

    Claude Code 2.1.226 returns its complete ``models`` array in one response.
    The stream parser still gathers every valid initialize response so a future
    paged/multi-frame implementation is not silently truncated. The protocol
    currently specifies no continuation request, so this function never
    fabricates one.
    """
    models: list[ClaudeModel] = []
    seen: set[tuple[str, str | None]] = set()
    received_models = False
    first_party = False
    for line in stdout.splitlines():
        try:
            event = json.loads(line)
        except json.JSONDecodeError:
            continue
        if not isinstance(event, dict) or event.get("type") != "control_response":
            continue
        response = event.get("response")
        if not isinstance(response, dict) or response.get("subtype") != "success":
            continue
        payload = response.get("response")
        if not isinstance(payload, dict) or not isinstance(payload.get("models"), list):
            continue
        received_models = True
        account = payload.get("account")
        if isinstance(account, dict) and account.get("apiProvider") == "firstParty":
            first_party = True
        for raw_model in payload["models"]:
            model = _model_info(raw_model)
            if model is None or len(models) >= MAX_MODELS:
                continue
            key = (model.value, model.resolved_model)
            if key not in seen:
                seen.add(key)
                models.append(model)
    if not received_models:
        return ClaudeModelDiscovery(
            (), False, False, "protocol-error", "protocol-error"
        )
    if not first_party:
        return ClaudeModelDiscovery(
            tuple(models), False, False, "not-subscription", "not-first-party"
        )
    if not models:
        return ClaudeModelDiscovery(
            (), True, True, "connected-subscription", "empty-model-catalog"
        )
    return ClaudeModelDiscovery(tuple(models), True, True, "connected-subscription")


def discover_claude_subscription_models(
    *,
    claude_bin: str | None = None,
    timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
    environment: Mapping[str, str] | None = None,
) -> ClaudeModelDiscovery:
    """Return current subscription models using only Claude's initialize protocol.

    A successful native authentication-status command or a locally cached model
    list is deliberately insufficient: ``live`` is true only after a valid,
    first-party initialize response supplies the account-visible ``models``.
    """
    binary = (
        claude_bin
        or os.environ.get("HERMES_CLAUDE_NATIVE_BIN")
        or DEFAULT_CLAUDE_BIN
    )
    if timeout_seconds <= 0:
        return ClaudeModelDiscovery(
            (), False, False, "invalid-request", "invalid-timeout"
        )
    stdout, error_code = _run_initialize(
        _command(binary), _native_environment(environment), timeout_seconds
    )
    if error_code:
        state = "unavailable" if error_code == "spawn-failed" else "degraded"
        return ClaudeModelDiscovery((), False, False, state, error_code)
    return _parse_initialize_output(stdout or "")