from __future__ import annotations

import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Literal

SCHEMA_VERSION = "1.0"
ManifestStatus = Literal["staged", "committed", "aborted"]
LockedStorageMode = Literal["compacted_session", "partitioned_by_run"]


def utc_now_iso() -> str:
    """Return an ISO-8601 UTC timestamp with second precision."""
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


@dataclass(slots=True)
class TableEntry:
    path: str
    format: Literal["parquet"]
    schema_fingerprint: str
    n_rows: int | None
    n_files: int
    file_hashes: list[str] | None
    created_by_step: str | None
    created_at: str
    locked_storage_mode: LockedStorageMode | None = None
    storage_mode_policy: dict[str, int] | None = None
    mode_lock_reason: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "path": self.path,
            "format": self.format,
            "schema_fingerprint": self.schema_fingerprint,
            "n_rows": self.n_rows,
            "n_files": self.n_files,
            "file_hashes": self.file_hashes,
            "created_by_step": self.created_by_step,
            "created_at": self.created_at,
            "locked_storage_mode": self.locked_storage_mode,
            "storage_mode_policy": self.storage_mode_policy,
            "mode_lock_reason": self.mode_lock_reason,
        }

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> TableEntry:
        locked_storage_mode = payload.get("locked_storage_mode")
        if locked_storage_mode is not None:
            locked_storage_mode = str(locked_storage_mode)
            if locked_storage_mode not in {"compacted_session", "partitioned_by_run"}:
                raise ValueError(f"invalid table locked_storage_mode: {locked_storage_mode}")

        storage_mode_policy = payload.get("storage_mode_policy")
        if storage_mode_policy is not None:
            storage_mode_policy = {
                str(key): int(value) for key, value in dict(storage_mode_policy).items()
            }

        return cls(
            path=str(payload["path"]),
            format="parquet",
            schema_fingerprint=str(payload["schema_fingerprint"]),
            n_rows=payload.get("n_rows"),
            n_files=int(payload["n_files"]),
            file_hashes=[str(item) for item in payload["file_hashes"]] if payload.get("file_hashes") else None,
            created_by_step=str(payload["created_by_step"]) if payload.get("created_by_step") is not None else None,
            created_at=str(payload["created_at"]),
            locked_storage_mode=locked_storage_mode,  # type: ignore[arg-type]
            storage_mode_policy=storage_mode_policy,
            mode_lock_reason=(
                str(payload["mode_lock_reason"]) if payload.get("mode_lock_reason") is not None else None
            ),
        )


@dataclass(slots=True)
class ArtifactEntry:
    path: str
    media_type: str | None
    sha256: str
    size_bytes: int
    created_by_step: str | None
    created_at: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "path": self.path,
            "media_type": self.media_type,
            "sha256": self.sha256,
            "size_bytes": self.size_bytes,
            "created_by_step": self.created_by_step,
            "created_at": self.created_at,
        }

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> ArtifactEntry:
        return cls(
            path=str(payload["path"]),
            media_type=str(payload["media_type"]) if payload.get("media_type") is not None else None,
            sha256=str(payload["sha256"]),
            size_bytes=int(payload["size_bytes"]),
            created_by_step=str(payload["created_by_step"]) if payload.get("created_by_step") is not None else None,
            created_at=str(payload["created_at"]),
        )


@dataclass(slots=True)
class RunManifest:
    pipeline_id: str
    run_id: str
    session_id: str | None = None
    session_run_index: int | None = None
    pipeline_label: str | None = None
    created_at: str = field(default_factory=utc_now_iso)
    status: ManifestStatus = "staged"
    metadata: dict[str, Any] | None = None
    tables: dict[str, TableEntry] = field(default_factory=dict)
    artifacts: dict[str, ArtifactEntry] = field(default_factory=dict)
    semantic_refs: dict[str, str] | None = None
    schema_version: str = SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema_version": self.schema_version,
            "pipeline_id": self.pipeline_id,
            "pipeline_label": self.pipeline_label,
            "run_id": self.run_id,
            "session_id": self.session_id,
            "session_run_index": self.session_run_index,
            "created_at": self.created_at,
            "status": self.status,
            "metadata": self.metadata,
            "tables": {name: entry.to_dict() for name, entry in self.tables.items()},
            "artifacts": {name: entry.to_dict() for name, entry in self.artifacts.items()},
            "semantic_refs": self.semantic_refs,
        }

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> RunManifest:
        status = str(payload.get("status", "staged"))
        if status not in {"staged", "committed", "aborted"}:
            raise ValueError(f"invalid manifest status: {status}")

        tables = {
            str(name): TableEntry.from_dict(entry)
            for name, entry in dict(payload.get("tables", {})).items()
        }
        artifacts = {
            str(name): ArtifactEntry.from_dict(entry)
            for name, entry in dict(payload.get("artifacts", {})).items()
        }

        semantic_refs = payload.get("semantic_refs")
        if semantic_refs is not None:
            semantic_refs = {str(key): str(value) for key, value in dict(semantic_refs).items()}

        metadata = payload.get("metadata")
        if metadata is not None:
            metadata = dict(metadata)

        return cls(
            schema_version=str(payload.get("schema_version", SCHEMA_VERSION)),
            pipeline_id=str(payload["pipeline_id"]),
            pipeline_label=str(payload["pipeline_label"]) if payload.get("pipeline_label") is not None else None,
            run_id=str(payload["run_id"]),
            session_id=str(payload["session_id"]) if payload.get("session_id") is not None else None,
            session_run_index=(
                int(payload["session_run_index"]) if payload.get("session_run_index") is not None else None
            ),
            created_at=str(payload["created_at"]),
            status=status,  # type: ignore[arg-type]
            metadata=metadata,
            tables=tables,
            artifacts=artifacts,
            semantic_refs=semantic_refs,
        )


@dataclass(slots=True)
class RunSummary:
    pipeline_id: str
    run_id: str
    created_at: str
    status: ManifestStatus
    session_id: str | None = None
    session_run_index: int | None = None
    pipeline_label: str | None = None


def manifest_to_json(manifest: RunManifest) -> str:
    return json.dumps(manifest.to_dict(), indent=2, sort_keys=True)


def manifest_from_json(payload: str) -> RunManifest:
    return RunManifest.from_dict(json.loads(payload))
