"""Atomic evidence writers using only the standard library and NumPy."""

from __future__ import annotations

import json
import os
import shutil
import tempfile
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from .hashing import file_sha256, to_jsonable
from .schema import SCHEMA_VERSION, STORAGE_FORMATS, ContractError


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _flush_and_sync(stream) -> None:
    stream.flush()
    os.fsync(stream.fileno())


def atomic_write_json(path: Path, document: Mapping[str, Any]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
    )
    temporary = Path(temporary_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as stream:
            json.dump(
                to_jsonable(document),
                stream,
                indent=2,
                sort_keys=True,
                ensure_ascii=False,
                allow_nan=False,
            )
            stream.write("\n")
            _flush_and_sync(stream)
        os.replace(temporary, path)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def atomic_write_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
    )
    temporary = Path(temporary_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as stream:
            for row in rows:
                stream.write(
                    json.dumps(
                        to_jsonable(row),
                        sort_keys=True,
                        separators=(",", ":"),
                        ensure_ascii=False,
                        allow_nan=False,
                    )
                )
                stream.write("\n")
            _flush_and_sync(stream)
        os.replace(temporary, path)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def atomic_write_npz(path: Path, arrays: Mapping[str, np.ndarray]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.parent / f".{path.stem}.{uuid.uuid4().hex}.tmp.npz"
    try:
        np.savez_compressed(temporary, **arrays)
        with temporary.open("rb") as stream:
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def normalize_sample_arrays(
    samples: Mapping[str, Any],
) -> tuple[dict[str, np.ndarray], int]:
    if not isinstance(samples, Mapping) or not samples:
        raise ContractError("trial samples must be a non-empty mapping")
    arrays: dict[str, np.ndarray] = {}
    sample_count: int | None = None
    for name, value in samples.items():
        if not isinstance(name, str) or not name:
            raise ContractError("sample field names must be non-empty strings")
        array = np.asarray(value)
        if array.dtype == object:
            raise ContractError(f"sample field {name!r} cannot have object dtype")
        if array.ndim == 0:
            raise ContractError(f"sample field {name!r} must have a sample axis")
        if sample_count is None:
            sample_count = int(array.shape[0])
        elif array.shape[0] != sample_count:
            raise ContractError(
                f"sample field {name!r} has {array.shape[0]} rows, "
                f"expected {sample_count}"
            )
        arrays[name] = array
    assert sample_count is not None
    if sample_count <= 0:
        raise ContractError("trial samples must contain at least one row")
    return arrays, sample_count


@dataclass(frozen=True)
class TrialPayload:
    samples: Mapping[str, Any]
    events: Sequence[Mapping[str, Any]] = field(default_factory=tuple)
    metadata: Mapping[str, Any] = field(default_factory=dict)


class AtomicTrialWriter:
    """Commit one successful trial directory using a final atomic rename."""

    def __init__(self, batch_dir: Path):
        self.batch_dir = Path(batch_dir)
        self.raw_dir = self.batch_dir / "raw"
        self.raw_dir.mkdir(parents=True, exist_ok=True)

    def trial_dir(self, trial_id: str) -> Path:
        return self.raw_dir / trial_id

    def write(self, trial: Mapping[str, Any], payload: TrialPayload) -> Path:
        trial_id = str(trial["trial_id"])
        destination = self.trial_dir(trial_id)
        if destination.exists():
            raise FileExistsError(f"trial artifact already exists: {destination}")
        arrays, sample_count = normalize_sample_arrays(payload.samples)
        temporary = self.raw_dir / f".{trial_id}.tmp.{uuid.uuid4().hex}"
        temporary.mkdir(parents=False, exist_ok=False)
        try:
            sample_path = temporary / "samples.npz"
            event_path = temporary / "events.jsonl"
            atomic_write_npz(sample_path, arrays)
            atomic_write_jsonl(event_path, list(payload.events))
            fields = {
                name: {"dtype": str(array.dtype), "shape": list(array.shape)}
                for name, array in sorted(arrays.items())
            }
            manifest = {
                "kind": "trial_manifest",
                "schema_version": SCHEMA_VERSION,
                "trial_id": trial_id,
                "pair_id": trial["pair_id"],
                "trial_spec_hash": trial["trial_spec_hash"],
                "status": "completed",
                "completed_utc": utc_now(),
                "sample_count": sample_count,
                "sample_fields": fields,
                "storage_formats": {
                    "samples": STORAGE_FORMATS["samples"],
                    "events": STORAGE_FORMATS["events"],
                },
                "files": {
                    "samples.npz": file_sha256(sample_path),
                    "events.jsonl": file_sha256(event_path),
                },
                "metadata": to_jsonable(payload.metadata),
            }
            atomic_write_json(temporary / "trial_manifest.json", manifest)
            os.replace(temporary, destination)
        except Exception:
            shutil.rmtree(temporary, ignore_errors=True)
            raise
        return destination


def record_trial_failure(
    batch_dir: Path,
    trial: Mapping[str, Any],
    error: BaseException,
) -> Path:
    failure_dir = Path(batch_dir) / "failures" / str(trial["trial_id"])
    failure_dir.mkdir(parents=True, exist_ok=True)
    path = failure_dir / f"attempt-{utc_now().replace(':', '')}-{uuid.uuid4().hex}.json"
    atomic_write_json(
        path,
        {
            "kind": "trial_failure",
            "schema_version": SCHEMA_VERSION,
            "trial_id": trial["trial_id"],
            "pair_id": trial["pair_id"],
            "trial_spec_hash": trial["trial_spec_hash"],
            "recorded_utc": utc_now(),
            "error_type": type(error).__name__,
            "error_message": str(error),
        },
    )
    return path