"""
Loads and validates config.yaml for the local LLM router.

Usage:
    from config import load_config
    cfg = load_config("config.yaml")
    cfg.objective.quality_tolerance  # etc.
"""

from __future__ import annotations

from pathlib import Path
from typing import Optional
from urllib.parse import urlparse

import yaml
from pydantic import BaseModel, ConfigDict, field_validator, model_validator


class StrictModel(BaseModel):
    """Base for every config section: an unknown key is an error.

    Pydantic ignores extra keys by default, which makes a typo or a
    misplaced setting silently do nothing while the file still loads and
    still looks configured. That is not hypothetical here — `max_input_chars`
    was written into the `verification:` block instead of `classifier:`,
    where it was accepted, ignored, and had no effect. It happened to carry
    the same value as the code default, so nothing visibly broke; editing it
    would simply have done nothing.

    Anyone tuning this file needs a wrong key to say so.
    """

    model_config = ConfigDict(extra="forbid")


class Objective(StrictModel):
    """What the router optimizes: quality, bounded by cost.

    Replaced a three-way weighted blend. See config.yaml for why — briefly,
    the cost weight was measured to be nearly inert while consuming 40% of
    every decision.
    """

    quality_tolerance: float = 0.10
    assumed_cache_rate: float = 0.917
    assumed_completion_tokens: int = 500
    max_energy_per_request: Optional[float] = None
    plan_kwh_per_period: Optional[float] = None

    @field_validator("quality_tolerance")
    @classmethod
    def tolerance_in_range(cls, v: float) -> float:
        if not (0.0 <= v < 1.0):
            raise ValueError("objective.quality_tolerance must be in [0, 1)")
        return v

    @field_validator("max_energy_per_request")
    @classmethod
    def ceiling_positive(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and v <= 0:
            raise ValueError(
                "objective.max_energy_per_request must be > 0 kWh, or null to disable"
            )
        return v


class ContextOverride(StrictModel):
    """Per-model context handling, for a row whose real limits are known.

    Typed rather than a bare ``dict`` so a typo INSIDE an override is an error
    too. That is the whole point of StrictModel, and it was not true here: the
    block validated, nothing read it, and config.yaml shipped a worked example
    for it -- so anyone who followed that example got silence.

    Both fields are optional; whichever is absent falls back to the global.
    """

    safety_factor: Optional[float] = None
    output_reserve_tokens: Optional[int] = None

    @field_validator("safety_factor")
    @classmethod
    def factor_in_range(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and not (0.0 < v <= 1.0):
            raise ValueError(
                "context.per_model_overrides[...].safety_factor must be in (0, 1]"
            )
        return v

    @field_validator("output_reserve_tokens")
    @classmethod
    def reserve_not_negative(cls, v: Optional[int]) -> Optional[int]:
        if v is not None and v < 0:
            raise ValueError(
                "context.per_model_overrides[...].output_reserve_tokens must be >= 0"
            )
        return v


class ContextConfig(StrictModel):
    safety_factor: float
    default_output_reserve_tokens: int
    # Read by poller.ModelRow.effective_context_window.
    per_model_overrides: dict[str, ContextOverride] = {}

    @field_validator("safety_factor")
    @classmethod
    def factor_in_range(cls, v: float) -> float:
        if not (0.0 < v <= 1.0):
            raise ValueError("context.safety_factor must be in (0, 1]")
        return v


class ProficiencyConfig(StrictModel):
    self_eval_min_samples: int
    leaderboard_weight: float
    self_eval_weight: float
    categories: list[str]

    @model_validator(mode="after")
    def blend_weights_sum_to_one(self) -> "ProficiencyConfig":
        total = round(self.leaderboard_weight + self.self_eval_weight, 6)
        if total != 1.0:
            raise ValueError(
                "proficiency.leaderboard_weight + self_eval_weight must sum to 1.0, "
                f"got {total}"
            )
        return self


class TieringConfig(StrictModel):
    cheap_completion_max: float
    tier1_context_max: float = float("inf")
    model_tiers: dict[str, int]

    @field_validator("cheap_completion_max")
    @classmethod
    def max_must_be_positive(cls, v: float) -> float:
        if v <= 0:
            raise ValueError("tiering.cheap_completion_max must be > 0")
        return v

    @field_validator("tier1_context_max")
    @classmethod
    def context_max_must_be_positive(cls, v: float) -> float:
        if v <= 0:
            raise ValueError("tiering.tier1_context_max must be > 0")
        return v

    @field_validator("model_tiers")
    @classmethod
    def overrides_in_range(cls, v: dict[str, int]) -> dict[str, int]:
        for model_id, tier in v.items():
            if tier not in (1, 2, 3):
                raise ValueError(
                    f"tiering.model_tiers[{model_id!r}] must be in {{1, 2, 3}}, "
                    f"got {tier}"
                )
        return v


class RoutingConfig(StrictModel):
    allowed_access_levels: list[str]
    default_latency_tolerance: str
    # Applied only when the REQUEST carries tool definitions. None disables it.
    min_tool_proficiency: Optional[float] = 0.5
    tool_use_category: str = "tool_use_agentic"

    # Whether a request carrying image parts is hard-restricted to vision-capable
    # models. A wrong guess here is a guaranteed 400, so this gates by default
    # and routing fails closed when the catalog flag is unknown.
    require_vision: bool = True
    # Whether a response_format requiring json_object/json_schema is hard-restricted
    # to JSON-mode-capable models. Same guaranteed-failure argument.
    require_json_mode: bool = True

    @field_validator("allowed_access_levels")
    @classmethod
    def levels_known(cls, v: list[str]) -> list[str]:
        known = {"public", "preview", "canary"}
        unknown = set(v) - known
        if unknown:
            raise ValueError(
                f"routing.allowed_access_levels contains unknown levels {sorted(unknown)}; "
                f"must be a subset of {sorted(known)}"
            )
        if not v:
            raise ValueError("routing.allowed_access_levels must not be empty")
        return v

    @field_validator("default_latency_tolerance")
    @classmethod
    def tolerance_known(cls, v: str) -> str:
        if v not in ("interactive", "batch"):
            raise ValueError(
                f"routing.default_latency_tolerance must be 'interactive' or 'batch', got {v!r}"
            )
        return v


class VerificationConfig(StrictModel):
    local_llm_enabled: bool = True
    min_completion_tokens: int = 600
    timeout_seconds: int = 60
    max_output_tokens: int = 1024
    outcome_attribution_window_seconds: int = 120
    # The local checker's OWN endpoint, no longer derived from the
    # classifier's. It speaks Ollama's NATIVE API (/api/chat, think=False),
    # which no cloud provider offers, so it must keep pointing at an Ollama
    # instance even when classification has been moved off this machine.
    base_url: str = "http://localhost:11434"
    # None means "whatever the classifier uses", which is correct only while
    # both run on the same local Ollama. Set it explicitly once they diverge.
    model: Optional[str] = None


class LocalVisionConfig(StrictModel):
    enabled: bool = True
    base_url: str = "http://localhost:11434/v1"   # OpenAI-compatible (classifier shape)
    api_key_env: Optional[str] = None
    model: str = "qwen3-vl:4b"
    timeout_seconds: int = 60
    max_images: int = 4
    max_image_bytes: int = 9 * 1024 * 1024        # 9 MiB, Ollama default cap

    @field_validator("timeout_seconds")
    @classmethod
    def timeout_positive(cls, v: int) -> int:
        if v <= 0:
            raise ValueError("local_vision.timeout_seconds must be > 0")
        return v

    @field_validator("max_images")
    @classmethod
    def images_positive(cls, v: int) -> int:
        if v <= 0:
            raise ValueError("local_vision.max_images must be > 0")
        return v

    @field_validator("max_image_bytes")
    @classmethod
    def image_bytes_positive(cls, v: int) -> int:
        if v <= 0:
            raise ValueError("local_vision.max_image_bytes must be > 0")
        return v


class EscalationConfig(StrictModel):
    enabled: bool
    max_tier: int
    min_confidence_before_bump: float
    # Off by default: the iteration budget escalates on evidence instead.
    preemptive_on_low_confidence: bool = False


class IterationConfig(StrictModel):
    """A tier's budget for corrective attempts after a verification failure."""

    enabled: bool = True
    attempts_by_tier: dict[int, int] = {1: 0, 2: 1, 3: 2}
    max_attempts_interactive: int = 1

    @field_validator("attempts_by_tier")
    @classmethod
    def attempts_sane(cls, v: dict[int, int]) -> dict[int, int]:
        for tier, attempts in v.items():
            if attempts < 0:
                raise ValueError(f"iteration.attempts_by_tier[{tier}] must be >= 0")
            if attempts > 5:
                raise ValueError(
                    f"iteration.attempts_by_tier[{tier}]={attempts} is implausibly "
                    "high; each attempt spends energy against a fixed quota"
                )
        return v


class FreshnessConfig(StrictModel):
    stale_after_days: int
    exclude_stale: bool
    exclude_deprecated: bool


class PinchConfig(StrictModel):
    """Optional relevance-based context pruning (Port of llmrouter's pinch).

    Prunes the provider-bound conversation — not the classifier input — when it
    exceeds ``budget_tokens``, so a long agent session ships fewer prompt tokens
    upstream. User/assistant/system messages are always kept; only tool results
    are summarized or dropped (they carry the bulk of a long session's tokens).
    """

    enabled: bool = False
    budget_tokens: int = 50000
    # How many recent user turns (plus their assistant replies and tool results)
    # are protected from pruning.
    keep_last_turns: int = 4
    # Tool results longer than this many characters are summarized in place.
    max_summarize_chars: int = 4000

    @field_validator("budget_tokens")
    @classmethod
    def budget_positive(cls, v: int) -> int:
        if v <= 0:
            raise ValueError("pinch.budget_tokens must be > 0")
        return v

    @field_validator("keep_last_turns")
    @classmethod
    def turns_positive(cls, v: int) -> int:
        if v <= 0:
            raise ValueError("pinch.keep_last_turns must be > 0")
        return v

    @field_validator("max_summarize_chars")
    @classmethod
    def summarize_chars_valid(cls, v: int) -> int:
        if v < 3000:
            raise ValueError(
                "pinch.max_summarize_chars must be >= 3000 (below this, "
                "summarization grows the message)"
            )
        return v


class DatabaseConfig(StrictModel):
    path: str


class ClassifierConfig(StrictModel):
    provider: str
    base_url: str
    # Env var holding the API key, for a classifier served by a provider that
    # actually checks one. None means unauthenticated, which is the local
    # Ollama case — it ignores the key entirely but the SDK requires one.
    api_key_env: Optional[str] = None
    model: str
    # Ceiling on the text handed to the classifier. 0 disables clamping.
    max_input_chars: int = 8000
    # When a preceding turn is available as context, frame the classifier
    # input as llmrouter does — "Context: <prev>\n---\nMessage: <task>" — so a
    # short follow-up ("Yes", "Try now?") can inherit the complexity of the
    # turn it continues instead of being classified in isolation as trivial.
    context_framing: bool = True
    timeout_seconds: int
    temperature: float = 0.0
    max_output_tokens: int = 1024
    fallback_tier: int = 2
    fallback_category: str = "general_chat"
    response_format: str
    system_prompt: str


class DispatchProvider(StrictModel):
    base_url: str
    api_key_env: str


class LoggingConfig(StrictModel):
    # log_path is gone. Nothing ever wrote a file: the dispatcher logs to
    # stderr and systemd captures that to the journal, so the setting named a
    # destination that did not exist.
    log_energy_observations: bool
    # Whether to write a row to route_decisions for every routing decision
    # (kind route | dispatch | chat | passthrough | local_vision). Off means
    # the monitoring TUI's decision history is empty; it does not affect
    # routing itself.
    log_route_decisions: bool = True
    # LLM_ROUTER_LOG_LEVEL overrides this at runtime — see logs.resolve_level.
    level: str = "info"

    @field_validator("level")
    @classmethod
    def level_known(cls, v: str) -> str:
        known = ("debug", "info", "warning", "error")
        if v.strip().lower() not in known:
            raise ValueError(f"logging.level must be one of {known}, got {v!r}")
        return v.strip().lower()


class RouterConfig(StrictModel):
    objective: Objective
    context: ContextConfig
    tiers: dict[int, str]
    tiering: TieringConfig
    proficiency: ProficiencyConfig
    routing: RoutingConfig
    verification: VerificationConfig = VerificationConfig()
    local_vision: LocalVisionConfig = LocalVisionConfig()
    escalation: EscalationConfig
    iteration: IterationConfig = IterationConfig()
    pinch: PinchConfig = PinchConfig()
    freshness: FreshnessConfig
    database: DatabaseConfig
    classifier: ClassifierConfig
    dispatch_providers: dict[str, DispatchProvider]
    logging: LoggingConfig

    @model_validator(mode="after")
    def tool_use_category_is_a_real_category(self) -> "RouterConfig":
        """The tool filter joins on a category name; a typo would disable it.

        A name that matches nothing produces NULL for every row, and NULL
        means "unproven, do not disqualify" — so the filter would silently
        pass everything. Failing at load beats a guard that quietly stops
        guarding.
        """
        if self.routing.min_tool_proficiency is None:
            return self
        if self.routing.tool_use_category not in self.proficiency.categories:
            raise ValueError(
                f"routing.tool_use_category "
                f"({self.routing.tool_use_category!r}) is not in "
                f"proficiency.categories — the tool-competence filter would "
                f"join against nothing and silently pass every model."
            )
        return self

    @model_validator(mode="after")
    def verifier_model_is_stated_once_the_hosts_differ(self) -> "RouterConfig":
        """A remote classifier must not lend its model name to the verifier.

        ``verification.model`` falling back to ``classifier.model`` is correct
        only while both point at the same Ollama. Once classification moves
        off-host the fallback names a model the local Ollama has never heard
        of, and the failure is SILENT: the verifier 404s, catches it, logs
        "local verification unavailable" and records no sample. Verification
        would appear to be on while producing nothing.

        Caught by pointing the classifier at NeuralWatt and watching the
        verifier POST ``deepseek-v4-flash`` to localhost:11434. Failing at
        config load instead means the misconfiguration is impossible rather
        than merely documented.
        """
        if not self.verification.local_llm_enabled or self.verification.model:
            return self
        classifier_host = urlparse(self.classifier.base_url).hostname
        verifier_host = urlparse(self.verification.base_url).hostname
        if classifier_host != verifier_host:
            raise ValueError(
                "verification.model must be set explicitly when the classifier "
                f"runs on a different host ({classifier_host} vs "
                f"{verifier_host}). It would otherwise fall back to "
                f"classifier.model ({self.classifier.model!r}), which the local "
                "Ollama does not serve — and the verifier fails silently. "
                "Set verification.model, or verification.local_llm_enabled: false."
            )
        return self


def load_config(path: str | Path = "config.yaml") -> RouterConfig:
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Config file not found: {path}")
    raw = yaml.safe_load(path.read_text())
    return RouterConfig(**raw)


def summary_lines(cfg: RouterConfig) -> list[str]:
    """What ``python config.py`` prints.

    A function rather than inline prints so a test can pin the attribute names.
    The previous version read ``cfg.weights``, which had been replaced by
    ``cfg.objective`` -- so the setup step documented in both README.md and
    CLAUDE.md said "Config loaded OK" and then died with AttributeError on a
    config that had in fact loaded perfectly.
    """
    return [
        "Config loaded OK",
        f"  objective: quality_tolerance={cfg.objective.quality_tolerance}, "
        f"max_energy_per_request={cfg.objective.max_energy_per_request}, "
        f"plan_kwh_per_period={cfg.objective.plan_kwh_per_period}",
        f"  categories: {cfg.proficiency.categories}",
        f"  classifier: {cfg.classifier.model} @ {cfg.classifier.base_url}",
        f"  verifier: {cfg.verification.model or cfg.classifier.model} "
        f"@ {cfg.verification.base_url}"
        + ("" if cfg.verification.local_llm_enabled else " (disabled)"),
        f"  tool filter: min_tool_proficiency={cfg.routing.min_tool_proficiency}",
        f"  dispatch providers: {list(cfg.dispatch_providers)}",
    ]


if __name__ == "__main__":
    import sys

    cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
    print("\n".join(summary_lines(load_config(cfg_path))))