from __future__ import annotations

import json
from pathlib import Path

from . import __version__
from .atif.models import (
    AgentSchema,
    FinalMetricsSchema,
    MetricsSchema,
    ObservationResultSchema,
    ObservationSchema,
    StepSchema,
    SubagentTrajectoryRefSchema,
    ToolCallSchema as AtifToolCallSchema,
    TrajectorySchema,
)
from .models import (
    GraderSpec,
    GraderResult,
    JudgeSpec,
    Outcome,
    OutcomePointer,
    PromptSpec,
    RunConfigFile,
    RunResult,
    RunSummary,
    SuiteConfig,
    Task,
    ToolCall,
    TranscriptEvent,
    TrialConfig,
    TrialSpec,
    TrialResult,
)


SCHEMA_MODELS = {
    "ToolCall": ToolCall,
    "OutcomePointer": OutcomePointer,
    "TranscriptEvent": TranscriptEvent,
    "Outcome": Outcome,
    "Task": Task,
    "PromptSpec": PromptSpec,
    "JudgeSpec": JudgeSpec,
    "GraderSpec": GraderSpec,
    "TrialSpec": TrialSpec,
    "TrialConfig": TrialConfig,
    "GraderResult": GraderResult,
    "TrialResult": TrialResult,
    "SuiteConfig": SuiteConfig,
    "RunConfigFile": RunConfigFile,
    "RunResult": RunResult,
    "RunSummary": RunSummary,
    "ATIF_Trajectory": TrajectorySchema,
    "ATIF_Step": StepSchema,
    "ATIF_ToolCall": AtifToolCallSchema,
    "ATIF_Observation": ObservationSchema,
    "ATIF_ObservationResult": ObservationResultSchema,
    "ATIF_SubagentTrajectoryRef": SubagentTrajectoryRefSchema,
    "ATIF_Metrics": MetricsSchema,
    "ATIF_FinalMetrics": FinalMetricsSchema,
    "ATIF_Agent": AgentSchema,
}


def export_json_schemas(out_dir: Path) -> list[Path]:
    out_dir.mkdir(parents=True, exist_ok=True)
    written: list[Path] = []
    for name, model in SCHEMA_MODELS.items():
        schema = model.model_json_schema()
        schema["$comment"] = f"generated by sentient-evals {__version__}"
        p = out_dir / f"{name}.schema.json"
        p.write_text(json.dumps(schema, indent=2, sort_keys=True))
        written.append(p)
    return written

