from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any, Literal

SCHEMA_VERSION = 5
RunStatus = Literal["running", "completed", "failed", "interrupted"]
LayerSummaryType = Literal["activation", "gradient"]
HookDirection = Literal["forward", "backward"]
ParameterRole = Literal["weight", "bias", "other"]
ParameterDistributionType = Literal["value", "gradient", "update"]
EventSeverity = Literal["debug", "info", "warning", "error", "critical"]
AlertSeverity = Literal["info", "warning", "error", "critical"]
AlertCategory = Literal["optimization", "representation", "generalization", "runtime"]
HealthScoreKey = Literal["overall", "optimization", "representation", "generalization", "runtime"]
TelemetryBoundary = Literal["step", "time"]
PhaseName = Literal[
    "initialization_adaptation",
    "rapid_learning",
    "stabilization",
    "stagnation",
    "overfitting_onset",
    "divergence",
]
ArtifactType = Literal[
    "histogram",
    "tensor_summary",
    "embedding_snapshot",
    "batch_diagnostics_snapshot",
    "checkpoint_metadata",
    "similarity_analysis_output",
    "alert_evidence_bundle",
]
ArtifactPayloadFormat = Literal["json", "text", "bytes"]

ACTIVE_RUN_STATUS: RunStatus = "running"
FINAL_RUN_STATUSES = {"completed", "failed", "interrupted"}


@dataclass(frozen=True, slots=True)
class ProjectRecord:
    id: str
    name: str
    root_path: str
    created_at: str
    updated_at: str
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class NoteRecord:
    id: str
    run_id: str
    note: str
    created_at: str
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class LifecycleEventRecord:
    id: str
    run_id: str
    event_type: str
    timestamp: str
    status: str
    payload: dict[str, Any]
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class EventRecord:
    id: str
    run_id: str
    sequence_id: int
    event_type: str
    severity: EventSeverity
    timestamp: str
    step: int | None = None
    epoch: int | None = None
    status: RunStatus | None = None
    payload: dict[str, Any] = field(default_factory=dict)
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class MetricRecord:
    id: str
    run_id: str
    sequence_id: int
    namespace: str
    metric_name: str
    metric_key: str
    value: float
    step: int | None
    epoch: int | None
    timestamp: str
    description: str | None = None
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class SystemTelemetryRecord:
    id: str
    run_id: str
    sequence_id: int
    boundary: TelemetryBoundary
    timestamp: str
    cpu_backend: str
    gpu_backend: str
    overflow_count: int
    skipped_step_count: int
    step: int | None = None
    epoch: int | None = None
    step_time_ms: float | None = None
    data_time_ms: float | None = None
    compute_time_ms: float | None = None
    cpu_usage_percent: float | None = None
    cpu_fallback_reason: str | None = None
    ram_used_bytes: int | None = None
    ram_total_bytes: int | None = None
    ram_available_bytes: int | None = None
    ram_usage_percent: float | None = None
    process_rss_bytes: int | None = None
    gpu_index: int | None = None
    gpu_name: str | None = None
    gpu_utilization_percent: float | None = None
    gpu_memory_used_bytes: int | None = None
    gpu_memory_total_bytes: int | None = None
    gpu_memory_usage_percent: float | None = None
    gpu_fallback_reason: str | None = None
    overflow_occurred: bool = False
    step_skipped: bool = False
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class BatchDiagnosticRecord:
    id: str
    run_id: str
    sequence_id: int
    loss: float
    timestamp: str
    step: int | None = None
    epoch: int | None = None
    batch_size: int | None = None
    sample_ids: tuple[str, ...] = ()
    class_counts: dict[str, int] = field(default_factory=dict)
    sequence_lengths: tuple[int, ...] = ()
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["sample_ids"] = list(self.sample_ids)
        payload["sequence_lengths"] = list(self.sequence_lengths)
        return payload


@dataclass(frozen=True, slots=True)
class LayerSummaryRecord:
    id: str
    run_id: str
    sequence_id: int
    step: int
    epoch: int | None
    module_name: str
    module_type: str
    summary_type: LayerSummaryType
    hook_direction: HookDirection
    call_index: int
    element_count: int
    timestamp: str
    tensor_shape: tuple[int, ...] | None = None
    mean: float | None = None
    std: float | None = None
    norm: float | None = None
    variance: float | None = None
    dead_unit_ratio: float | None = None
    saturation_ratio: float | None = None
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["tensor_shape"] = list(self.tensor_shape) if self.tensor_shape is not None else None
        return payload


@dataclass(frozen=True, slots=True)
class ParameterSummaryRecord:
    id: str
    run_id: str
    sequence_id: int
    step: int
    epoch: int | None
    timestamp: str
    module_name: str
    module_type: str | None
    parameter_name: str
    parameter_key: str
    parameter_role: ParameterRole
    element_count: int
    tensor_shape: tuple[int, ...] | None = None
    dtype: str | None = None
    requires_grad: bool | None = None
    mean: float | None = None
    std: float | None = None
    variance: float | None = None
    norm: float | None = None
    min: float | None = None
    max: float | None = None
    abs_mean: float | None = None
    abs_max: float | None = None
    zero_ratio: float | None = None
    near_zero_ratio: float | None = None
    finite_ratio: float | None = None
    nan_count: int | None = None
    inf_count: int | None = None
    grad_mean: float | None = None
    grad_std: float | None = None
    grad_norm: float | None = None
    grad_abs_max: float | None = None
    update_norm: float | None = None
    update_mean: float | None = None
    update_std: float | None = None
    update_to_parameter_norm_ratio: float | None = None
    optimizer_name: str | None = None
    optimizer_group_index: int | None = None
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["tensor_shape"] = list(self.tensor_shape) if self.tensor_shape is not None else None
        return payload


@dataclass(frozen=True, slots=True)
class ParameterDistributionRecord:
    id: str
    run_id: str
    sequence_id: int
    step: int
    epoch: int | None
    timestamp: str
    module_name: str
    module_type: str | None
    parameter_name: str
    parameter_key: str
    parameter_role: ParameterRole
    distribution_type: ParameterDistributionType
    bin_edges: tuple[float, ...]
    counts: tuple[int, ...]
    sample_count: int
    density: bool = False
    total_count: int | None = None
    min: float | None = None
    max: float | None = None
    nan_count: int | None = None
    inf_count: int | None = None
    clipped_low_count: int | None = None
    clipped_high_count: int | None = None
    summary_id: str | None = None
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["bin_edges"] = list(self.bin_edges)
        payload["counts"] = list(self.counts)
        return payload


@dataclass(frozen=True, slots=True)
class ArtifactRecord:
    id: str
    run_id: str
    sequence_id: int
    artifact_type: ArtifactType
    namespace: str
    artifact_name: str
    artifact_key: str
    payload_format: ArtifactPayloadFormat
    content_type: str
    payload_path: str
    payload_relative_path: str
    metadata_path: str
    metadata_relative_path: str
    step: int | None
    epoch: int | None
    timestamp: str
    size_bytes: int
    sha256: str
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class EvidenceRecord:
    kind: str
    source: str
    summary: str
    observed_value: Any = None
    threshold_value: Any = None
    step: int | None = None
    step_start: int | None = None
    step_end: int | None = None
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class AlertRecord:
    id: str
    run_id: str
    sequence_id: int
    rule_id: str
    category: AlertCategory
    severity: AlertSeverity
    title: str
    summary: str
    fact: str
    inference: str
    timestamp: str
    step: int | None = None
    step_start: int | None = None
    step_end: int | None = None
    evidence: tuple[EvidenceRecord, ...] = field(default_factory=tuple)
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["evidence"] = [evidence.to_dict() for evidence in self.evidence]
        return payload


@dataclass(frozen=True, slots=True)
class HealthScoreRecord:
    id: str
    run_id: str
    sequence_id: int
    score_key: HealthScoreKey
    category: str
    score: float
    status: str
    summary: str
    timestamp: str
    step_start: int | None = None
    step_end: int | None = None
    evidence: tuple[EvidenceRecord, ...] = field(default_factory=tuple)
    contributing_alert_ids: tuple[str, ...] = field(default_factory=tuple)
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["evidence"] = [evidence.to_dict() for evidence in self.evidence]
        payload["contributing_alert_ids"] = list(self.contributing_alert_ids)
        return payload


@dataclass(frozen=True, slots=True)
class PhaseRecord:
    id: str
    run_id: str
    sequence_id: int
    phase_name: PhaseName
    display_name: str
    summary: str
    step_start: int
    step_end: int
    timestamp: str
    evidence: tuple[EvidenceRecord, ...] = field(default_factory=tuple)
    metadata: dict[str, Any] | None = None
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["evidence"] = [evidence.to_dict() for evidence in self.evidence]
        return payload


@dataclass(frozen=True, slots=True)
class RunRecord:
    id: str
    project_id: str
    project_name: str
    status: RunStatus
    created_at: str
    started_at: str
    updated_at: str
    ended_at: str | None
    parent_run_id: str | None
    lineage_root_run_id: str
    resume_count: int
    latest_resume_at: str | None
    config_path: str
    environment_path: str
    run_path: str
    git_commit: str | None
    strict_mode: bool
    tags: tuple[str, ...] = ()
    notes: tuple[NoteRecord, ...] = ()
    lifecycle_events: tuple[LifecycleEventRecord, ...] = ()
    schema_version: int = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["notes"] = [note.to_dict() for note in self.notes]
        payload["lifecycle_events"] = [event.to_dict() for event in self.lifecycle_events]
        return payload


@dataclass(frozen=True, slots=True)
class RunContext:
    storage_dir: str
    store_root: str
    db_path: str
    project: ProjectRecord
    run: RunRecord

    @property
    def project_id(self) -> str:
        return self.project.id

    @property
    def project_name(self) -> str:
        return self.project.name

    @property
    def run_id(self) -> str:
        return self.run.id

    @property
    def run_path(self) -> str:
        return self.run.run_path

    @property
    def status(self) -> RunStatus:
        return self.run.status

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema_version": SCHEMA_VERSION,
            "storage_dir": self.storage_dir,
            "store_root": self.store_root,
            "db_path": self.db_path,
            "project": self.project.to_dict(),
            "run": self.run.to_dict(),
        }
