"""Data models for the alice CLI."""

from __future__ import annotations

import re
from dataclasses import dataclass
from enum import Enum
from typing import Any

from pydantic import BaseModel, field_validator

# AWS region pattern: 2-letter continent code + direction/name + digit(s)
# e.g. us-east-1, eu-west-2, ap-southeast-1, ca-central-1, sa-east-1
_AWS_REGION_RE = re.compile(r"^[a-z]{2}-[a-z]+-\d+$")

# ARN pattern: arn:<partition>:<service>:<region>:<account>:<resource>
_ARN_RE = re.compile(r"^arn:[a-z][a-z0-9\-]*:[a-z0-9\-]+:[a-z0-9\-]*:\d*:[a-zA-Z0-9\-_/:.]+")


class CLIConfig(BaseModel):
    """Resolved CLI configuration with defaults and overrides.

    Uses Pydantic for runtime validation of configuration values.
    """

    profile: str | None = None
    region: str
    namespace: str
    environment: str
    quiet: bool
    debug: bool = False
    personality: str = "alice"
    bedrock_access_key: str | None = None
    bedrock_secret_key: str | None = None
    bedrock_bearer_token: str | None = None
    inference_profile_arns: dict[str, str] | None = None

    model_config = {"frozen": False, "arbitrary_types_allowed": True}

    @field_validator("region")
    @classmethod
    def validate_region(cls, v: str) -> str:
        """Validate AWS region format.

        Enforces the standard AWS region pattern (e.g. us-east-1, eu-west-2).
        Custom/local regions that don't match will raise ValueError.
        """
        if not v:
            raise ValueError("Region cannot be empty")
        if not _AWS_REGION_RE.match(v):
            raise ValueError(
                f"Invalid AWS region format: {v!r}. "
                "Expected format: <continent>-<direction>-<number> (e.g. us-east-1)"
            )
        return v

    @field_validator("namespace")
    @classmethod
    def validate_namespace(cls, v: str) -> str:
        """Validate namespace is non-empty."""
        if not v:
            raise ValueError("Namespace cannot be empty")
        return v

    @field_validator("environment")
    @classmethod
    def validate_environment(cls, v: str) -> str:
        """Validate environment is non-empty."""
        if not v:
            raise ValueError("Environment cannot be empty")
        return v

    @field_validator("personality")
    @classmethod
    def validate_personality(cls, v: str) -> str:
        """Validate personality is one of the allowed values."""
        allowed = {"alice", "mr_rogers", "shodan"}
        if v not in allowed:
            raise ValueError(f"Personality must be one of {allowed}, got '{v}'")
        return v

    @field_validator("inference_profile_arns")
    @classmethod
    def validate_inference_profile_arns(cls, v: dict[str, str] | None) -> dict[str, str] | None:
        """Validate that all inference profile ARNs have valid ARN format."""
        if v is None:
            return v
        for alias, arn in v.items():
            if arn and not _ARN_RE.match(arn):
                raise ValueError(
                    f"Invalid ARN format for alias {alias!r}: {arn!r}. "
                    "Expected format: arn:<partition>:<service>:<region>:<account>:<resource>"
                )
        return v


@dataclass
class UsernameResult:
    """Result of JHED detection from an STS caller identity ARN."""

    username: str
    session_name: str
    arn: str
    account_id: str = ""
    is_federated: bool = False  # True when role is AWSReservedSSO_* (org/federation account)


@dataclass
class CredentialSet:
    """Bedrock credentials retrieved from Secrets Manager."""

    bearer_token: str
    access_key: str
    secret_key: str
    region: str
    status: str
    raw_json: str
    inference_profile_arns: dict[str, str]  # model alias → AIP ARN


class OutputFormat(str, Enum):
    """Credential output format."""

    EXPORT = "export"
    EVAL = "eval"
    DOTENV = "dotenv"
    JSON = "json"


class SaveFormat(str, Enum):
    """Session save output format."""

    MARKDOWN = "markdown"
    JSON = "json"
    TEXT = "text"


@dataclass(frozen=True)
class ModelAlias:
    """A friendly alias for a Bedrock model ID."""

    alias: str
    model_id: str
    description: str


# Curated model aliases — short names matching the inference_profile_model_sources
# keys in terraform/envs/drcc-ai/terraform.tfvars.  The model IDs here are the
# cross-region inference profile IDs (us.*) so they work with the Converse API.
MODEL_ALIASES: list[ModelAlias] = [
    # ── Anthropic Claude ──────────────────────────────────────────────
    ModelAlias(
        "claude-sonnet-4.6", "us.anthropic.claude-sonnet-4-6", "Claude Sonnet 4.6 (default)"
    ),
    ModelAlias("claude-opus-4.6", "us.anthropic.claude-opus-4-6-v1", "Claude Opus 4.6"),
    ModelAlias(
        "claude-sonnet-4.5", "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "Claude Sonnet 4.5"
    ),
    ModelAlias("claude-opus-4.5", "us.anthropic.claude-opus-4-5-20251101-v1:0", "Claude Opus 4.5"),
    ModelAlias(
        "claude-haiku-4.5", "us.anthropic.claude-haiku-4-5-20251001-v1:0", "Claude Haiku 4.5"
    ),
    # ── Amazon Nova ───────────────────────────────────────────────────
    ModelAlias("nova-pro", "amazon.nova-pro-v1:0", "Amazon Nova Pro"),
    ModelAlias("nova-lite", "amazon.nova-lite-v1:0", "Amazon Nova Lite"),
    ModelAlias("nova-micro", "amazon.nova-micro-v1:0", "Amazon Nova Micro"),
    # ── Mistral AI ────────────────────────────────────────────────────
    ModelAlias(
        "mistral-large-3", "mistral.mistral-large-3-675b-instruct", "Mistral Large 3 (675B)"
    ),
    ModelAlias("magistral-small", "mistral.magistral-small-2509", "Magistral Small"),
    ModelAlias("devstral-2", "mistral.devstral-2-123b", "Devstral 2 (123B)"),
    ModelAlias("ministral-14b", "mistral.ministral-3-14b-instruct", "Ministral 14B"),
    ModelAlias("ministral-8b", "mistral.ministral-3-8b-instruct", "Ministral 8B"),
    ModelAlias("ministral-3b", "mistral.ministral-3-3b-instruct", "Ministral 3B"),
    # ── DeepSeek ──────────────────────────────────────────────────────
    ModelAlias("deepseek-r1", "us.deepseek.r1-v1:0", "DeepSeek R1"),
    # ── Meta Llama ────────────────────────────────────────────────────
    ModelAlias(
        "llama4-maverick", "us.meta.llama4-maverick-17b-instruct-v1:0", "Llama 4 Maverick 17B"
    ),
    ModelAlias("llama4-scout", "us.meta.llama4-scout-17b-instruct-v1:0", "Llama 4 Scout 17B"),
    # ── Convenience shorthands ────────────────────────────────────────
    ModelAlias("sonnet", "us.anthropic.claude-sonnet-4-6", "→ claude-sonnet-4.6"),
    ModelAlias("opus", "us.anthropic.claude-opus-4-6-v1", "→ claude-opus-4.6"),
    ModelAlias("haiku", "us.anthropic.claude-haiku-4-5-20251001-v1:0", "→ claude-haiku-4.5"),
]

# Quick lookup: alias → ModelAlias
_ALIAS_MAP: dict[str, ModelAlias] = {m.alias: m for m in MODEL_ALIASES}

DEFAULT_MODEL_ALIAS = "sonnet"
DEFAULT_MODEL_ID = _ALIAS_MAP[DEFAULT_MODEL_ALIAS].model_id

# Claude Code default model environment variables.
# Maps env-var name → (alias used for resolve_model_id, fallback model ID).
CLAUDE_DEFAULT_MODELS: dict[str, tuple[str, str]] = {
    "ANTHROPIC_DEFAULT_OPUS_MODEL": ("opus", _ALIAS_MAP["opus"].model_id),
    "ANTHROPIC_DEFAULT_SONNET_MODEL": ("sonnet", _ALIAS_MAP["sonnet"].model_id),
    "ANTHROPIC_DEFAULT_HAIKU_MODEL": ("haiku", _ALIAS_MAP["haiku"].model_id),
}


_CONTEXT_SUFFIX_RE = re.compile(r":\d+[kKmMgG]$")


def _strip_context_suffix(model_id: str) -> str:
    """Remove Bedrock list-models context-window suffixes (e.g. ':300k').

    ``alice list-models`` surfaces raw Bedrock model IDs that include a
    context-size qualifier appended by the listing API, such as
    ``amazon.nova-pro-v1:0:300k``.  The Converse API only accepts the base
    model ID (``amazon.nova-pro-v1:0``), so the suffix must be stripped
    before use.  ARNs and IDs that don't match the pattern are returned
    unchanged.
    """
    return _CONTEXT_SUFFIX_RE.sub("", model_id)


def resolve_model_id(value: str, inference_profile_arns: dict[str, str] | None = None) -> str:
    """Resolve a model alias to an inference profile ARN or model ID.

    Resolution order:
    1. If *value* matches a known alias and the user has an inference profile
       ARN for that alias (or its canonical name), return the ARN.
    2. If *value* matches a known alias but no ARN is available, return the
       alias's registered model ID (direct invocation).
    3. Otherwise assume *value* is already a full model ID / ARN and return
       it with any context-window suffix stripped.
    """
    profiles = inference_profile_arns or {}
    lower = value.lower()
    hit = _ALIAS_MAP.get(lower)
    if hit is None:
        # Not a known alias — pass through, but strip any context-size suffix
        # (e.g. "amazon.nova-pro-v1:0:300k" → "amazon.nova-pro-v1:0")
        return _strip_context_suffix(value)

    # Determine the canonical alias name (e.g. "sonnet" → "claude-sonnet-4.6")
    canonical = hit.alias
    # For convenience shorthands, find the canonical alias that shares the same model_id
    if hit.description.startswith("→"):
        for m in MODEL_ALIASES:
            if m.model_id == hit.model_id and not m.description.startswith("→"):
                canonical = m.alias
                break

    # Prefer the user's inference profile ARN if available
    arn = profiles.get(canonical) or profiles.get(hit.alias)
    if arn:
        return arn

    return hit.model_id


def model_help_text() -> str:
    """Build a help string listing available aliases."""
    # Only show the primary aliases (skip convenience shorthands that start with →)
    primary = [m for m in MODEL_ALIASES if not m.description.startswith("→")]
    aliases = ", ".join(m.alias for m in primary)
    return (
        f"Model alias or full Bedrock model ID [default: sonnet]. "
        f"Aliases: {aliases}. "
        f"Run 'alice list-aliases' for details."
    )


@dataclass
class TokenUsage:
    """Token counts from a Converse API response."""

    input: int
    output: int
    total: int


@dataclass
class SessionRecord:
    """Unified session record for all command types."""

    id: str  # UUID4 string
    type: str  # "chat" | "invoke" | "dialog" | "batch" | "compare"
    timestamp: str  # ISO 8601 UTC string
    model: str  # Model alias or ID used
    prompt: str  # Original prompt text
    response: str | None  # Model response (None for multi-turn chat)
    tokens: TokenUsage | None  # Token usage (None if unavailable)
    metadata: dict[str, Any] | None  # Arbitrary extra data


@dataclass
class IndexEntry:
    """Lightweight metadata entry in the LockerIndex manifest."""

    id: str  # Session UUID
    type: str  # Session type
    timestamp: str  # ISO 8601 UTC
    model: str  # Model alias/ID
    prompt_preview: str  # First 80 chars of prompt
    tokens: TokenUsage | None


@dataclass
class LockerIndex:
    """The index manifest — a list of IndexEntry objects."""

    entries: list[IndexEntry]


@dataclass
class ChatMessage:
    """A single message in a chat conversation."""

    role: str  # "user" | "assistant"
    content: str


class SummaryLength(str, Enum):
    """Summary verbosity level."""

    SHORT = "short"
    MEDIUM = "medium"
    DETAILED = "detailed"


SUMMARY_PROMPTS: dict[SummaryLength, str] = {
    SummaryLength.SHORT: "Summarize the following in 2-3 sentences:\n\n{text}",
    SummaryLength.MEDIUM: "Summarize the following in 1-2 paragraphs:\n\n{text}",
    SummaryLength.DETAILED: (
        "Provide a comprehensive summary preserving key details and structure:\n\n{text}"
    ),
}


@dataclass
class DoctorCheck:
    """A single diagnostic check result."""

    name: str
    passed: bool
    detail: str
    hint: str  # Remediation hint shown on failure
