from __future__ import annotations

import json
import re
import threading
from collections.abc import Callable, Mapping, Sequence
from numbers import Integral, Real
from pathlib import Path
from typing import Any, ClassVar

from models import ArtifactRecord, RunContext
from persistence.artifacts import ARTIFACT_STORAGE_LAYOUT, ensure_artifact_layout
from storage import LocalSessionStore, utc_now

_IDENTITY_SEGMENT_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")
_SUPPORTED_ARTIFACT_TYPES = {
    "histogram",
    "tensor_summary",
    "embedding_snapshot",
    "batch_diagnostics_snapshot",
    "checkpoint_metadata",
    "similarity_analysis_output",
    "alert_evidence_bundle",
}
_SNAPSHOT_ARTIFACT_TYPES = {
    "embedding_snapshot",
    "batch_diagnostics_snapshot",
    "checkpoint_metadata",
}
_ARTIFACT_DIRECTORY_BY_TYPE = {
    "histogram": "histograms",
    "tensor_summary": "summaries",
    "embedding_snapshot": "embeddings",
    "batch_diagnostics_snapshot": "snapshots",
    "checkpoint_metadata": "checkpoints",
    "similarity_analysis_output": "analysis",
    "alert_evidence_bundle": "evidence",
}
_CONTENT_TYPE_BY_FORMAT = {
    "json": "application/json",
    "text": "text/plain; charset=utf-8",
    "bytes": "application/octet-stream",
}
_EXTENSION_BY_FORMAT = {
    "json": "json",
    "text": "txt",
    "bytes": "bin",
}

RetentionHook = Callable[["ArtifactStore", ArtifactRecord], None]


class ArtifactStore:
    """Run-scoped persistence for non-scalar artifact payloads and indexed metadata."""

    _sequence_lock: ClassVar[threading.RLock] = threading.RLock()
    _next_sequence_by_run: ClassVar[dict[tuple[str, str], int]] = {}

    def __init__(
        self,
        session: Any,
        *,
        retention_hooks: Sequence[RetentionHook] | None = None,
        raise_on_retention_error: bool | None = None,
    ) -> None:
        self._context, self._store = self._resolve_session_dependencies(session)
        self._retention_hooks = self._normalize_retention_hooks(retention_hooks)
        if raise_on_retention_error is not None and not isinstance(raise_on_retention_error, bool):
            raise TypeError("raise_on_retention_error must be a boolean or None.")
        self._raise_on_retention_error = (
            raise_on_retention_error
            if raise_on_retention_error is not None
            else self._context.run.strict_mode
        )
        self._closed = False

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

    def store(
        self,
        artifact_type: str,
        name: str,
        payload: Any,
        *,
        step: int | None = None,
        epoch: int | None = None,
        namespace: str | None = None,
        payload_format: str | None = None,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> ArtifactRecord:
        """Persist an artifact payload plus indexed metadata for the current run."""

        self._ensure_open()
        normalized_artifact_type = self._normalize_artifact_type(artifact_type)
        normalized_step = self._normalize_index("step", step)
        normalized_epoch = self._normalize_index("epoch", epoch)
        normalized_namespace, artifact_name, artifact_key, artifact_segments = (
            self._normalize_artifact_identity(name, namespace=namespace)
        )
        normalized_metadata = self._normalize_metadata(metadata)
        normalized_payload_format = self._normalize_payload_format(payload_format, payload)
        record = self._store.persist_artifact(
            self.run_id,
            artifact_type=normalized_artifact_type,
            namespace=normalized_namespace,
            artifact_name=artifact_name,
            artifact_key=artifact_key,
            payload=payload,
            payload_format=normalized_payload_format,
            metadata=normalized_metadata,
            step=normalized_step,
            epoch=normalized_epoch,
            timestamp=timestamp or utc_now(),
        )
        self._run_retention_hooks(record)
        return record

    def log_histogram(
        self,
        name: str,
        *,
        bins: Sequence[float],
        counts: Sequence[float],
        step: int | None = None,
        epoch: int | None = None,
        namespace: str | None = None,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> ArtifactRecord:
        """Persist a histogram artifact using a stable JSON payload schema."""

        normalized_bins = self._normalize_numeric_sequence(bins, label="bins")
        normalized_counts = self._normalize_numeric_sequence(
            counts,
            label="counts",
            allow_negative=False,
        )
        if len(normalized_bins) < 2:
            raise ValueError("Histogram bins must include at least two edges.")
        if len(normalized_counts) != len(normalized_bins) - 1:
            raise ValueError("Histogram counts must contain exactly len(bins) - 1 values.")
        if any(left >= right for left, right in zip(normalized_bins, normalized_bins[1:])):
            raise ValueError("Histogram bins must be strictly increasing.")

        return self.store(
            "histogram",
            name,
            {
                "bins": normalized_bins,
                "counts": normalized_counts,
            },
            step=step,
            epoch=epoch,
            namespace=namespace,
            metadata=metadata,
            timestamp=timestamp,
        )

    def log_snapshot(
        self,
        name: str,
        payload: Any,
        *,
        snapshot_type: str = "batch_diagnostics_snapshot",
        step: int | None = None,
        epoch: int | None = None,
        namespace: str | None = None,
        payload_format: str | None = None,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> ArtifactRecord:
        """Persist a snapshot-style artifact such as a batch or checkpoint snapshot."""

        normalized_snapshot_type = self._normalize_artifact_type(snapshot_type)
        if normalized_snapshot_type not in _SNAPSHOT_ARTIFACT_TYPES:
            raise ValueError(
                "snapshot_type must be one of "
                f"{sorted(_SNAPSHOT_ARTIFACT_TYPES)}."
            )

        return self.store(
            normalized_snapshot_type,
            name,
            payload,
            step=step,
            epoch=epoch,
            namespace=namespace,
            payload_format=payload_format,
            metadata=metadata,
            timestamp=timestamp,
        )

    def read(
        self,
        *,
        artifact_id: str | None = None,
        artifact_type: str | None = None,
        artifact_types: Sequence[str] | None = None,
        namespace: str | None = None,
        artifact_name: str | None = None,
        artifact_key: str | None = None,
        payload_format: str | None = None,
        step: int | None = None,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
        limit: int | None = None,
    ) -> tuple[ArtifactRecord, ...]:
        """Return persisted artifact metadata for the run in sequence order."""

        normalized_artifact_id = self._normalize_optional_text(artifact_id, label="artifact_id")
        normalized_artifact_type = (
            self._normalize_artifact_type(artifact_type) if artifact_type is not None else None
        )
        normalized_artifact_types = (
            tuple(self._normalize_artifact_type(item) for item in artifact_types)
            if artifact_types is not None
            else None
        )
        normalized_namespace = (
            self._normalize_namespace(namespace) if namespace is not None else None
        )
        normalized_artifact_name = (
            self._normalize_segment(artifact_name, "Artifact name")
            if artifact_name is not None
            else None
        )
        normalized_artifact_key = (
            self._normalize_key_filter(artifact_key) if artifact_key is not None else None
        )
        normalized_payload_format = (
            self._normalize_payload_format_label(payload_format)
            if payload_format is not None
            else None
        )

        return self._store.read_artifacts(
            self.run_id,
            artifact_id=normalized_artifact_id,
            artifact_type=normalized_artifact_type,
            artifact_types=normalized_artifact_types,
            namespace=normalized_namespace,
            artifact_name=normalized_artifact_name,
            artifact_key=normalized_artifact_key,
            payload_format=normalized_payload_format,
            step=self._normalize_index("step", step),
            step_min=self._normalize_index("step_min", step_min),
            step_max=self._normalize_index("step_max", step_max),
            epoch=self._normalize_index("epoch", epoch),
            epoch_min=self._normalize_index("epoch_min", epoch_min),
            epoch_max=self._normalize_index("epoch_max", epoch_max),
            limit=limit,
        )

    def load(self, artifact: str | ArtifactRecord) -> Any:
        """Load and deserialize an artifact payload from its persisted file."""

        record = artifact
        if isinstance(artifact, str):
            records = self.read(artifact_id=artifact, limit=1)
            if not records:
                raise LookupError(f"Artifact '{artifact}' does not exist in run '{self.run_id}'.")
            record = records[0]
        if not isinstance(record, ArtifactRecord):
            raise TypeError("artifact must be an ArtifactRecord or artifact ID string.")

        return self._store.load_artifact_payload(self.run_id, record.id)

    def close(self) -> None:
        self._closed = True

    def __enter__(self) -> ArtifactStore:
        return self

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        self.close()

    def _allocate_sequence_id(self) -> int:
        run_key = (self._context.db_path, self.run_id)
        with self._sequence_lock:
            next_sequence = self._next_sequence_by_run.get(run_key)
            if next_sequence is None:
                next_sequence = self._store.get_next_artifact_sequence(self.run_id)
            self._next_sequence_by_run[run_key] = next_sequence + 1
            return next_sequence

    def _ensure_layout_definition(self) -> None:
        ensure_artifact_layout(Path(self._context.run_path))

    def _resolve_storage_paths(
        self,
        *,
        artifact_type: str,
        artifact_segments: tuple[str, ...],
        payload_format: str,
        sequence_id: int,
        step: int | None,
        epoch: int | None,
    ) -> tuple[Path, str, Path, str]:
        run_path = Path(self._context.run_path)
        artifacts_root = run_path / "artifacts" / _ARTIFACT_DIRECTORY_BY_TYPE[artifact_type]
        artifact_directory = artifacts_root.joinpath(*artifact_segments)
        file_name = (
            f"step-{self._format_index_label(step)}"
            f"__epoch-{self._format_index_label(epoch)}"
            f"__seq-{sequence_id:08d}.{_EXTENSION_BY_FORMAT[payload_format]}"
        )
        payload_path = artifact_directory / file_name
        metadata_path = payload_path.with_suffix(".meta.json")
        return (
            payload_path,
            payload_path.relative_to(run_path).as_posix(),
            metadata_path,
            metadata_path.relative_to(run_path).as_posix(),
        )

    @staticmethod
    def _write_payload_file(path: Path, payload: bytes) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        temp_path = path.with_suffix(f"{path.suffix}.tmp")
        temp_path.write_bytes(payload)
        temp_path.replace(path)

    @staticmethod
    def _cleanup_file(path: Path) -> None:
        try:
            path.unlink(missing_ok=True)
        except Exception:
            pass

    def _run_retention_hooks(self, record: ArtifactRecord) -> None:
        for hook in self._retention_hooks:
            try:
                hook(self, record)
            except Exception as exc:
                if self._raise_on_retention_error:
                    raise RuntimeError(
                        f"Retention hook failed for artifact '{record.id}' in run '{self.run_id}'."
                    ) from exc

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("ArtifactStore is closed.")

    @staticmethod
    def _resolve_session_dependencies(session: Any) -> tuple[RunContext, LocalSessionStore]:
        if isinstance(session, RunContext):
            return session, LocalSessionStore(session.storage_dir)

        context = getattr(session, "context", None)
        store = getattr(session, "_store", None)
        if isinstance(context, RunContext) and isinstance(store, LocalSessionStore):
            return context, store

        raise TypeError(
            "ArtifactStore expects a TrainingSession or RunContext created by the Quark SDK."
        )

    @staticmethod
    def _normalize_retention_hooks(
        retention_hooks: Sequence[RetentionHook] | None,
    ) -> tuple[RetentionHook, ...]:
        if retention_hooks is None:
            return ()
        normalized_hooks = tuple(retention_hooks)
        if not all(callable(hook) for hook in normalized_hooks):
            raise TypeError("retention_hooks must contain only callable hook functions.")
        return normalized_hooks

    @staticmethod
    def _normalize_artifact_type(artifact_type: str) -> str:
        candidate = artifact_type.strip().lower()
        if candidate not in _SUPPORTED_ARTIFACT_TYPES:
            raise ValueError(
                f"artifact_type '{artifact_type}' is not supported. "
                f"Expected one of {sorted(_SUPPORTED_ARTIFACT_TYPES)}."
            )
        return candidate

    @staticmethod
    def _normalize_artifact_identity(
        name: str,
        *,
        namespace: str | None,
    ) -> tuple[str, str, str, tuple[str, ...]]:
        candidate = name.strip()
        if not candidate:
            raise ValueError("Artifact names must not be empty.")

        if namespace is not None:
            normalized_namespace = ArtifactStore._normalize_namespace(namespace)
            if "/" in candidate:
                raise ValueError(
                    "Artifact names must not include '/' when an explicit namespace is provided."
                )
            normalized_name = ArtifactStore._normalize_segment(candidate, "Artifact name")
        else:
            parts = candidate.split("/")
            if len(parts) == 1:
                normalized_namespace = ""
                normalized_name = ArtifactStore._normalize_segment(parts[0], "Artifact name")
            else:
                normalized_namespace = "/".join(
                    ArtifactStore._normalize_segment(part, "Namespace segment")
                    for part in parts[:-1]
                )
                normalized_name = ArtifactStore._normalize_segment(parts[-1], "Artifact name")

        artifact_key = (
            f"{normalized_namespace}/{normalized_name}"
            if normalized_namespace
            else normalized_name
        )
        segments = tuple(part for part in artifact_key.split("/") if part)
        return normalized_namespace, normalized_name, artifact_key, segments

    @staticmethod
    def _normalize_namespace(namespace: str) -> str:
        parts = [segment.strip() for segment in namespace.split("/")]
        if not parts or any(not part for part in parts):
            raise ValueError("Artifact namespaces must not be empty.")
        return "/".join(
            ArtifactStore._normalize_segment(part, "Namespace segment") for part in parts
        )

    @staticmethod
    def _normalize_segment(value: str, label: str) -> str:
        candidate = value.strip()
        if not candidate:
            raise ValueError(f"{label} must not be empty.")
        if _IDENTITY_SEGMENT_PATTERN.fullmatch(candidate) is None:
            raise ValueError(
                f"{label} '{candidate}' may only contain letters, digits, '.', '_' or '-'."
            )
        return candidate

    @staticmethod
    def _normalize_key_filter(artifact_key: str) -> str:
        parts = [segment.strip() for segment in artifact_key.split("/")]
        if not parts or any(not part for part in parts):
            raise ValueError("artifact_key must not be empty.")
        return "/".join(
            ArtifactStore._normalize_segment(part, "Artifact key segment") for part in parts
        )

    @staticmethod
    def _normalize_optional_text(value: str | None, *, label: str) -> str | None:
        if value is None:
            return None
        normalized = value.strip()
        if not normalized:
            raise ValueError(f"{label} must not be empty when provided.")
        return normalized

    @staticmethod
    def _normalize_index(label: str, value: int | None) -> int | None:
        if value is None:
            return None
        if isinstance(value, bool) or not isinstance(value, Integral):
            raise TypeError(f"{label} must be an integer when provided.")
        normalized_value = int(value)
        if normalized_value < 0:
            raise ValueError(f"{label} must not be negative.")
        return normalized_value

    @staticmethod
    def _normalize_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any] | None:
        if metadata is None:
            return None
        payload = dict(metadata)
        json.dumps(payload, sort_keys=True)
        return payload

    @staticmethod
    def _normalize_payload_format(payload_format: str | None, payload: Any) -> str:
        if payload_format is not None:
            return ArtifactStore._normalize_payload_format_label(payload_format)
        if isinstance(payload, str):
            return "text"
        if isinstance(payload, (bytes, bytearray, memoryview)):
            return "bytes"
        return "json"

    @staticmethod
    def _normalize_payload_format_label(payload_format: str) -> str:
        candidate = payload_format.strip().lower()
        if candidate not in _CONTENT_TYPE_BY_FORMAT:
            raise ValueError(
                f"payload_format '{payload_format}' is not supported. "
                f"Expected one of {sorted(_CONTENT_TYPE_BY_FORMAT)}."
            )
        return candidate

    @staticmethod
    def _serialize_payload(payload: Any, *, payload_format: str) -> bytes:
        if payload_format == "json":
            return json.dumps(payload, indent=2, sort_keys=True).encode("utf-8")
        if payload_format == "text":
            if not isinstance(payload, str):
                raise TypeError("text payloads must be Python strings.")
            return payload.encode("utf-8")
        if not isinstance(payload, (bytes, bytearray, memoryview)):
            raise TypeError("bytes payloads must be bytes-like objects.")
        return bytes(payload)

    @staticmethod
    def _normalize_numeric_sequence(
        values: Sequence[float],
        *,
        label: str,
        allow_negative: bool = True,
    ) -> list[float]:
        normalized_values: list[float] = []
        for raw_value in values:
            if isinstance(raw_value, bool) or not isinstance(raw_value, Real):
                raise TypeError(f"{label} values must be numeric.")
            candidate = float(raw_value)
            if candidate != candidate or candidate in {float("inf"), float("-inf")}:
                raise ValueError(f"{label} values must be finite real numbers.")
            if not allow_negative and candidate < 0:
                raise ValueError(f"{label} values must not be negative.")
            normalized_values.append(candidate)
        return normalized_values

    @staticmethod
    def _format_index_label(value: int | None) -> str:
        if value is None:
            return "na"
        return f"{value:08d}"
