"""Immutable, hash-addressed JSON artifacts for GAP runs."""

from __future__ import annotations

import hashlib
import json
import os
import re
import tempfile
from pathlib import Path
from typing import Any

from pydantic import BaseModel

from .models import ModelCallRecord, StageArtifact


def canonical_json(value: Any) -> str:
    if isinstance(value, BaseModel):
        value = value.model_dump(mode="json")
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    )


def sha256_payload(value: Any) -> str:
    return hashlib.sha256(canonical_json(value).encode("utf-8")).hexdigest()


def safe_name(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("._") or "artifact"


class RunStore:
    """Writes immutable artifacts and refuses conflicting overwrites."""

    def __init__(self, root: Path, item_id: str) -> None:
        self.root = root
        self.item_id = item_id
        self.item_root = root / "items" / safe_name(item_id)
        for directory in [
            self.item_root,
            self.item_root / "calls",
            self.item_root / "stages",
            self.item_root / "iterations",
        ]:
            directory.mkdir(parents=True, exist_ok=True)

    def write_json(self, relative_path: str | Path, value: Any) -> Path:
        path = self.item_root / relative_path
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = value.model_dump(mode="json") if isinstance(value, BaseModel) else value
        text = json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
        if path.exists():
            existing = path.read_text(encoding="utf-8")
            if existing != text:
                raise FileExistsError(
                    f"refusing to overwrite non-identical artifact {path}"
                )
            return path
        descriptor, temp_name = tempfile.mkstemp(
            prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
        )
        try:
            with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
                handle.write(text)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temp_name, path)
        finally:
            if os.path.exists(temp_name):
                os.unlink(temp_name)
        return path

    def write_input(self, value: Any) -> Path:
        return self.write_json("input.json", value)

    def write_config(self, value: Any) -> Path:
        return self.write_json("config.json", value)

    def write_call(self, record: ModelCallRecord) -> Path:
        return self.write_json(
            Path("calls") / f"{safe_name(record.request_id)}.json", record
        )

    def write_stage(
        self,
        stage: str,
        payload: BaseModel | dict[str, Any],
        *,
        request_id: str | None,
    ) -> Path:
        raw = payload.model_dump(mode="json") if isinstance(payload, BaseModel) else payload
        artifact = StageArtifact(
            item_id=self.item_id,
            stage=stage,
            payload_sha256=sha256_payload(raw),
            payload=raw,
            request_id=request_id,
        )
        return self.write_json(Path("stages") / f"{safe_name(stage)}.json", artifact)

    def write_iteration(self, iteration: int, value: Any) -> Path:
        return self.write_json(
            Path("iterations") / f"{iteration:02d}.json",
            value,
        )

    def write_final(self, value: Any) -> Path:
        return self.write_json("final.json", value)