from __future__ import annotations

from dataclasses import dataclass, fields, is_dataclass
from pathlib import Path
from typing import Callable, TypeVar

from .runtime import ExecutionRecord, StepTrace


StateT = TypeVar("StateT")
CandidateT = TypeVar("CandidateT")
MetadataT = TypeVar("MetadataT")


def to_jsonable(value: object) -> object:
    if is_dataclass(value):
        return {
            field.name: to_jsonable(getattr(value, field.name))
            for field in fields(value)
        }
    if isinstance(value, tuple):
        return [to_jsonable(item) for item in value]
    if isinstance(value, list):
        return [to_jsonable(item) for item in value]
    if isinstance(value, dict):
        return {str(key): to_jsonable(item) for key, item in value.items()}
    return value


@dataclass(frozen=True)
class SerializedStepTrace:
    previous_state: object
    next_state: object
    candidate: object
    accepted: bool
    elapsed_seconds: float
    metadata: object | None


@dataclass(frozen=True)
class SerializedExecutionRecord:
    accepted: tuple[SerializedStepTrace, ...]
    attempts: tuple[SerializedStepTrace, ...]
    final_state: object
    total_seconds: float


def serialize_step_trace(
    trace: StepTrace[StateT, CandidateT, MetadataT],
    *,
    state_encoder: Callable[[StateT], object] = to_jsonable,
    candidate_encoder: Callable[[CandidateT], object] = to_jsonable,
    metadata_encoder: Callable[[MetadataT | None], object | None] = to_jsonable,
) -> SerializedStepTrace:
    return SerializedStepTrace(
        previous_state=state_encoder(trace.previous_state),
        next_state=state_encoder(trace.next_state),
        candidate=candidate_encoder(trace.candidate),
        accepted=trace.accepted,
        elapsed_seconds=trace.elapsed_seconds,
        metadata=metadata_encoder(trace.metadata),
    )


def serialize_execution_record(
    record: ExecutionRecord[StateT, CandidateT, MetadataT],
    *,
    state_encoder: Callable[[StateT], object] = to_jsonable,
    candidate_encoder: Callable[[CandidateT], object] = to_jsonable,
    metadata_encoder: Callable[[MetadataT | None], object | None] = to_jsonable,
) -> SerializedExecutionRecord:
    return SerializedExecutionRecord(
        accepted=tuple(
            serialize_step_trace(
                trace,
                state_encoder=state_encoder,
                candidate_encoder=candidate_encoder,
                metadata_encoder=metadata_encoder,
            )
            for trace in record.accepted
        ),
        attempts=tuple(
            serialize_step_trace(
                trace,
                state_encoder=state_encoder,
                candidate_encoder=candidate_encoder,
                metadata_encoder=metadata_encoder,
            )
            for trace in record.attempts
        ),
        final_state=state_encoder(record.final_state),
        total_seconds=record.total_seconds,
    )


def deserialize_step_trace(
    data: SerializedStepTrace | dict[str, object],
    *,
    state_decoder: Callable[[object], StateT],
    candidate_decoder: Callable[[object], CandidateT],
    metadata_decoder: Callable[[object | None], MetadataT | None],
) -> StepTrace[StateT, CandidateT, MetadataT]:
    if isinstance(data, SerializedStepTrace):
        payload = data
    else:
        payload = SerializedStepTrace(
            previous_state=data["previous_state"],
            next_state=data["next_state"],
            candidate=data["candidate"],
            accepted=bool(data["accepted"]),
            elapsed_seconds=float(data["elapsed_seconds"]),
            metadata=data.get("metadata"),
        )
    return StepTrace(
        previous_state=state_decoder(payload.previous_state),
        next_state=state_decoder(payload.next_state),
        candidate=candidate_decoder(payload.candidate),
        accepted=payload.accepted,
        elapsed_seconds=payload.elapsed_seconds,
        metadata=metadata_decoder(payload.metadata),
    )


def deserialize_execution_record(
    data: SerializedExecutionRecord | dict[str, object],
    *,
    state_decoder: Callable[[object], StateT],
    candidate_decoder: Callable[[object], CandidateT],
    metadata_decoder: Callable[[object | None], MetadataT | None],
) -> ExecutionRecord[StateT, CandidateT, MetadataT]:
    if isinstance(data, SerializedExecutionRecord):
        payload = data
    else:
        payload = SerializedExecutionRecord(
            accepted=tuple(data["accepted"]),  # type: ignore[arg-type]
            attempts=tuple(data["attempts"]),  # type: ignore[arg-type]
            final_state=data["final_state"],
            total_seconds=float(data["total_seconds"]),
        )
    return ExecutionRecord(
        accepted=tuple(
            deserialize_step_trace(
                step,
                state_decoder=state_decoder,
                candidate_decoder=candidate_decoder,
                metadata_decoder=metadata_decoder,
            )
            for step in payload.accepted
        ),
        attempts=tuple(
            deserialize_step_trace(
                step,
                state_decoder=state_decoder,
                candidate_decoder=candidate_decoder,
                metadata_decoder=metadata_decoder,
            )
            for step in payload.attempts
        ),
        final_state=state_decoder(payload.final_state),
        total_seconds=payload.total_seconds,
    )


def save_execution_record_json(
    record: ExecutionRecord[StateT, CandidateT, MetadataT],
    path: str | Path,
    *,
    state_encoder: Callable[[StateT], object] = to_jsonable,
    candidate_encoder: Callable[[CandidateT], object] = to_jsonable,
    metadata_encoder: Callable[[MetadataT | None], object | None] = to_jsonable,
) -> None:
    from .artifacts import save_artifact_json

    serialized = serialize_execution_record(
        record,
        state_encoder=state_encoder,
        candidate_encoder=candidate_encoder,
        metadata_encoder=metadata_encoder,
    )
    save_artifact_json("execution_trace", serialized, path)
