"""Artifact registry — manages writing, hashing, and indexing for a run."""

from __future__ import annotations

import hashlib
import json
import logging
from datetime import UTC, datetime
from pathlib import Path
from collections.abc import Callable
from typing import Any

from dpe.engine.states import STAGE_ARTIFACT_MAP, Stage
from dpe.models.manifest import ArtifactRecord

logger = logging.getLogger(__name__)


class ArtifactRegistry:
    """Writes artifacts to disk and tracks them for manifest generation."""

    def __init__(
        self,
        run_dir: Path,
        clock: Callable[[], datetime] | None = None,
    ) -> None:
        self._run_dir = run_dir
        self._run_dir.mkdir(parents=True, exist_ok=True)
        self._artifacts: list[ArtifactRecord] = []
        self._clock = clock or (lambda: datetime.now(UTC))

    @property
    def artifacts(self) -> list[ArtifactRecord]:
        return list(self._artifacts)

    @property
    def run_dir(self) -> Path:
        return self._run_dir

    def write_stage_artifact(self, stage: Stage, content: str) -> ArtifactRecord:
        """Write an artifact for the given stage and return its record."""
        filename = STAGE_ARTIFACT_MAP[stage]
        return self._write(filename=filename, stage=stage.value, content=content)

    def write_manifest(self, manifest_data: dict[str, Any]) -> Path:
        """Write run_manifest.json to the run directory."""
        path = self._run_dir / "run_manifest.json"
        payload = json.dumps(manifest_data, indent=2, default=str)
        path.write_text(payload, encoding="utf-8")
        logger.info("Manifest written: %s", path)
        return path

    def write_artifact_index(self) -> Path:
        """Write artifact_index.json summarising all artifacts."""
        path = self._run_dir / "artifact_index.json"
        index = [a.model_dump(mode="json") for a in self._artifacts]
        path.write_text(json.dumps(index, indent=2, default=str), encoding="utf-8")
        logger.info("Artifact index written: %s (%d entries)", path, len(index))
        return path

    def _write(self, *, filename: str, stage: str, content: str) -> ArtifactRecord:
        """Write a file, compute its hash, and record it."""
        if "/" in filename or "\\" in filename:
            raise ValueError(f"Invalid artifact filename: {filename}")

        path = self._run_dir / filename
        encoded = content.encode("utf-8")
        path.write_bytes(encoded)

        sha = hashlib.sha256(encoded).hexdigest()
        record = ArtifactRecord(
            filename=filename,
            stage=stage,
            sha256=sha,
            size_bytes=len(encoded),
            written_at=self._clock(),
        )
        self._artifacts.append(record)
        logger.info("Artifact written: %s (%d bytes, sha256=%s…)", filename, len(encoded), sha[:12])
        return record
