from __future__ import annotations

from pathlib import Path
import argparse
import json

from .artifacts import (
    ARTIFACT_SCHEMA_VERSION,
    ArtifactManifest,
    ArtifactManifestEntry,
    save_manifest_json,
)
from .demo_registry import available_demo_names, get_demo_definition
from .graph import GraphEdgeSpec, GraphSchema, categorizer_node, critic_node, generator_node, policy_node
from .reporting import save_run_report_json, summarize_sequence_run
from .runtime import run_until_acceptance_count
from .serialization import save_execution_record_json
from .snapshots import create_demo_snapshot, resume_demo_snapshot, save_demo_snapshot_json


def _graph_schema_for_system(system: object) -> GraphSchema:
    graph_nodes = (
        generator_node("generator", system.generator),  # type: ignore[attr-defined]
        critic_node("critic", system.critic),  # type: ignore[attr-defined]
        categorizer_node("categorizer", system.categorizer),  # type: ignore[attr-defined]
        policy_node("policy", system.policy),  # type: ignore[attr-defined]
    )
    return GraphSchema(
        nodes=tuple(node.spec() for node in graph_nodes),
        edges=(
            GraphEdgeSpec("generator", "candidate", "critic", "candidate"),
            GraphEdgeSpec("generator", "candidate", "categorizer", "candidate"),
            GraphEdgeSpec("generator", "candidate", "policy", "candidate"),
            GraphEdgeSpec("critic", "critique", "policy", "critique"),
            GraphEdgeSpec("categorizer", "category", "policy", "category"),
        ),
    )


def export_demo_artifacts(
    output_dir: str | Path,
    *,
    demo_name: str,
    accepted_count: int = 2,
    max_attempts_per_accept: int = 4,
    snapshot_after_accepted: int | None = None,
) -> dict[str, Path]:
    definition = get_demo_definition(demo_name)
    destination = Path(output_dir)
    destination.mkdir(parents=True, exist_ok=True)
    system = definition.build_system()

    graph_path = destination / "graph.json"
    trace_path = destination / "trace.json"
    report_path = destination / "report.json"
    manifest_path = destination / "manifest.json"
    snapshot_path = destination / "snapshot.json"

    manifest_entries = [
        ArtifactManifestEntry("graph_schema", graph_path.name),
        ArtifactManifestEntry("execution_trace", trace_path.name),
        ArtifactManifestEntry("run_report", report_path.name),
    ]

    if snapshot_after_accepted is None:
        record = run_until_acceptance_count(
            system,
            definition.initial_state,
            accepted_count=accepted_count,
            max_attempts_per_accept=max_attempts_per_accept,
        )
    else:
        if snapshot_after_accepted < 0:
            raise ValueError("snapshot_after_accepted must be non-negative")
        if snapshot_after_accepted > accepted_count:
            raise ValueError("snapshot_after_accepted cannot exceed accepted_count")
        partial_record = run_until_acceptance_count(
            system,
            definition.initial_state,
            accepted_count=snapshot_after_accepted,
            max_attempts_per_accept=max_attempts_per_accept,
        )
        snapshot = create_demo_snapshot(
            demo_name,
            system=system,
            record=partial_record,
            parameters={
                "accepted_count": snapshot_after_accepted,
                "max_attempts_per_accept": max_attempts_per_accept,
            },
        )
        save_demo_snapshot_json(snapshot, snapshot_path)
        manifest_entries.append(ArtifactManifestEntry("demo_snapshot", snapshot_path.name))
        if snapshot_after_accepted == accepted_count:
            record = partial_record
        else:
            system, record = resume_demo_snapshot(
                snapshot,
                additional_accepted_count=accepted_count - snapshot_after_accepted,
                max_attempts_per_accept=max_attempts_per_accept,
            )

    _graph_schema_for_system(system).save_json(graph_path)
    save_execution_record_json(record, trace_path)
    report = summarize_sequence_run(
        record,
        sequence_getter=definition.sequence_getter,
        alphabet_size=definition.alphabet_size,
        parameters={
            "example": definition.name,
            "accepted_count": accepted_count,
            "max_attempts_per_accept": max_attempts_per_accept,
        },
    )
    save_run_report_json(report, report_path)
    save_manifest_json(
        ArtifactManifest(
            schema_version=ARTIFACT_SCHEMA_VERSION,
            artifacts=tuple(manifest_entries),
            metadata={
                "example": definition.name,
                "accepted_count": accepted_count,
                "max_attempts_per_accept": max_attempts_per_accept,
                "snapshot_after_accepted": snapshot_after_accepted,
            },
        ),
        manifest_path,
    )
    artifacts = {
        "graph": graph_path,
        "trace": trace_path,
        "report": report_path,
        "manifest": manifest_path,
    }
    if snapshot_after_accepted is not None:
        artifacts["snapshot"] = snapshot_path
    return artifacts


def export_xor_demo_artifacts(
    output_dir: str | Path,
    *,
    accepted_count: int = 2,
    max_attempts_per_accept: int = 4,
) -> dict[str, Path]:
    return export_demo_artifacts(
        output_dir,
        demo_name="xor_novelty",
        accepted_count=accepted_count,
        max_attempts_per_accept=max_attempts_per_accept,
    )


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="synaptopus-demo-export",
        description="Export graph, trace, and report artifacts for a Synaptopus internal demo.",
    )
    parser.add_argument("output_dir")
    parser.add_argument(
        "--demo",
        default="xor_novelty",
        choices=available_demo_names(),
    )
    parser.add_argument("--accepted-count", type=int, default=2)
    parser.add_argument("--max-attempts-per-accept", type=int, default=4)
    parser.add_argument("--snapshot-after-accepted", type=int)
    args = parser.parse_args(argv)

    artifacts = export_demo_artifacts(
        args.output_dir,
        demo_name=args.demo,
        accepted_count=args.accepted_count,
        max_attempts_per_accept=args.max_attempts_per_accept,
        snapshot_after_accepted=args.snapshot_after_accepted,
    )
    print(json.dumps({name: str(path) for name, path in artifacts.items()}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
