"""Deterministic JSON Schema generation from the canonical Python types."""

from __future__ import annotations

import json
from copy import deepcopy
from pathlib import Path
from typing import Any

from pydantic import BaseModel

from .models import (
    ActionSelectionOutput,
    BattleState,
    BenchmarkLogRecord,
    BestOfThreeResult,
    Decision,
    MirroredPairPlan,
    MirroredPairResult,
    SelectionVarianceResult,
    SystemOneChoiceAnswer,
    SystemOneRequest,
    TeamSelectionInput,
    TeamSelectionOutput,
    TournamentManifest,
    TournamentSummary,
)

SCHEMA_MODELS: dict[str, type[BaseModel]] = {
    "team_selection_input": TeamSelectionInput,
    "team_selection_output": TeamSelectionOutput,
    "battle_state": BattleState,
    "action_selection_output": ActionSelectionOutput,
    "system_one_request": SystemOneRequest,
    "system_one_choice_answer": SystemOneChoiceAnswer,
    "decision": Decision,
    "mirrored_pair_plan": MirroredPairPlan,
    "mirrored_pair_result": MirroredPairResult,
    "selection_variance_result": SelectionVarianceResult,
    "best_of_three_result": BestOfThreeResult,
    "benchmark_log_record": BenchmarkLogRecord,
    "tournament_manifest": TournamentManifest,
    "tournament_summary": TournamentSummary,
}


def schema_for(name: str) -> dict[str, Any]:
    """Generate one deterministic Draft 2020-12-compatible JSON Schema."""

    try:
        model = SCHEMA_MODELS[name]
    except KeyError as error:
        raise ValueError(f"unknown schema: {name}") from error
    return model.model_json_schema(mode="validation", ref_template="#/$defs/{model}")


def _remove_annotations(value: Any) -> Any:
    if isinstance(value, dict):
        return {
            key: _remove_annotations(item)
            for key, item in value.items()
            if key not in {"title", "default", "examples"}
        }
    if isinstance(value, list):
        return [_remove_annotations(item) for item in value]
    return value


def structured_output_schema(name: str) -> dict[str, Any]:
    """Return a compact schema suitable for strict model structured output."""

    return _remove_annotations(deepcopy(schema_for(name)))


def export_schemas(directory: Path) -> list[Path]:
    """Materialize generated schemas for inspection or another AI provider."""

    directory.mkdir(parents=True, exist_ok=True)
    paths = []
    for name in sorted(SCHEMA_MODELS):
        path = directory / f"{name.replace('_', '-')}.schema.json"
        path.write_text(
            json.dumps(schema_for(name), indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        paths.append(path)
    return paths