from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import json

from .artifacts import save_artifact_json
from .demo_registry import get_demo_definition
from .runtime import ExecutionRecord, merge_execution_records, run_until_acceptance_count
from .serialization import (
    SerializedExecutionRecord,
    deserialize_execution_record,
    serialize_execution_record,
)


@dataclass(frozen=True)
class DemoSnapshot:
    demo_name: str
    system: dict[str, object]
    record: SerializedExecutionRecord
    parameters: dict[str, object]


def create_demo_snapshot(
    demo_name: str,
    *,
    system: object,
    record: ExecutionRecord[object, object, object],
    parameters: dict[str, object] | None = None,
) -> DemoSnapshot:
    definition = get_demo_definition(demo_name)
    serialized_record = serialize_execution_record(
        record,
        state_encoder=definition.state_encoder,
    )
    return DemoSnapshot(
        demo_name=demo_name,
        system=definition.system_encoder(system),
        record=serialized_record,
        parameters=dict(parameters or {}),
    )


def save_demo_snapshot_json(snapshot: DemoSnapshot, path: str | Path) -> None:
    save_artifact_json(
        "demo_snapshot",
        snapshot,
        path,
        metadata={"demo_name": snapshot.demo_name},
    )


def load_demo_snapshot_json(path: str | Path) -> DemoSnapshot:
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    if payload["artifact_type"] != "demo_snapshot":
        raise ValueError("artifact is not a demo_snapshot")
    envelope_payload = payload["payload"]
    record_payload = envelope_payload["record"]
    return DemoSnapshot(
        demo_name=str(envelope_payload["demo_name"]),
        system=dict(envelope_payload["system"]),
        record=SerializedExecutionRecord(
            accepted=tuple(record_payload["accepted"]),
            attempts=tuple(record_payload["attempts"]),
            final_state=record_payload["final_state"],
            total_seconds=float(record_payload["total_seconds"]),
        ),
        parameters=dict(envelope_payload.get("parameters", {})),
    )


def restore_demo_snapshot(
    snapshot: DemoSnapshot,
) -> tuple[object, ExecutionRecord[object, object, object]]:
    definition = get_demo_definition(snapshot.demo_name)
    system = definition.system_decoder(snapshot.system)
    record = deserialize_execution_record(
        snapshot.record,
        state_decoder=definition.state_decoder,
        candidate_decoder=definition.candidate_decoder,
        metadata_decoder=definition.metadata_decoder,
    )
    return system, record


def resume_demo_snapshot(
    snapshot: DemoSnapshot,
    *,
    additional_accepted_count: int,
    max_attempts_per_accept: int,
) -> tuple[object, ExecutionRecord[object, object, object]]:
    if additional_accepted_count < 0:
        raise ValueError("additional_accepted_count must be non-negative")
    system, record = restore_demo_snapshot(snapshot)
    if additional_accepted_count == 0:
        return system, record
    continuation = run_until_acceptance_count(
        system,
        record.final_state,
        accepted_count=additional_accepted_count,
        max_attempts_per_accept=max_attempts_per_accept,
    )
    return system, merge_execution_records(record, continuation)
