from __future__ import annotations

import json
import os
import shutil
import stat
import threading
from contextlib import suppress
from pathlib import Path
from typing import Any, Final
from uuid import uuid4

from k1link.artifacts import write_json_atomic
from k1link.simulation.contracts import (
    ControlSetpoint,
    QualificationArtifact,
    QualificationEvent,
    QualificationRun,
    RunKind,
    RunState,
    SimulationContractError,
    control_setpoint_from_dict,
)

MAX_RECORD_BYTES: Final = 1024 * 1024
SEQUENCE_WIDTH: Final = 20
ACTIVE_RECOVERY_STATES: Final = {
    RunState.STARTING,
    RunState.RUNNING,
    RunState.PAUSED,
    RunState.STOPPING,
}
COMMANDABLE_RUN_KINDS: Final = {
    RunKind.SIMULATION_CLOSED_LOOP,
    RunKind.DIGITAL_TWIN_CLOSED_LOOP,
}
TRANSITIONS: Final[dict[RunState, frozenset[RunState]]] = {
    RunState.ADMITTED: frozenset({RunState.STARTING, RunState.ABORTED}),
    RunState.STARTING: frozenset({RunState.RUNNING, RunState.FAILED, RunState.ABORTED}),
    RunState.RUNNING: frozenset(
        {RunState.PAUSED, RunState.STOPPING, RunState.FAILED, RunState.ABORTED}
    ),
    RunState.PAUSED: frozenset(
        {RunState.RUNNING, RunState.STOPPING, RunState.FAILED, RunState.ABORTED}
    ),
    RunState.STOPPING: frozenset({RunState.COMPLETED, RunState.FAILED, RunState.ABORTED}),
    RunState.COMPLETED: frozenset(),
    RunState.FAILED: frozenset(),
    RunState.ABORTED: frozenset(),
}


class QualificationRunStoreError(RuntimeError):
    """Base error for the append-only qualification-run repository."""


class QualificationRunNotFoundError(QualificationRunStoreError):
    """The requested qualification run is not present."""


class QualificationRunConflictError(QualificationRunStoreError):
    """The caller used stale revision/sequence state or an existing identity."""


class QualificationRunIntegrityError(QualificationRunStoreError):
    """Stored qualification evidence violates its immutable contract."""


class QualificationRunTransitionError(QualificationRunStoreError):
    """A requested lifecycle transition is not allowed."""


class QualificationRunStore:
    """Single-writer, crash-safe, append-only Polygon run repository.

    The run manifest is immutable. Lifecycle, domain events, accepted commands
    and artifact-index entries are one-file-per-record journals created with
    exclusive filesystem semantics and fsynced before acknowledgement.
    """

    def __init__(self, root: Path, *, read_only: bool = False) -> None:
        configured_root = root.expanduser()
        if read_only:
            if not configured_root.is_dir():
                raise QualificationRunIntegrityError("qualification repository is missing")
            _reject_symlink(configured_root, "qualification repository")
            self.root = configured_root.resolve()
        else:
            self.root = configured_root.resolve()
            self.root.mkdir(mode=0o700, parents=True, exist_ok=True)
            _reject_symlink(self.root, "qualification repository")
            _chmod_private(self.root)
        self.read_only = read_only
        self._lock = threading.RLock()

    def create(self, run: QualificationRun) -> QualificationRun:
        self._require_writable()
        if (
            run.state is not RunState.ADMITTED
            or run.revision != 0
            or run.artifacts
            or run.started_at_utc is not None
            or run.ended_at_utc is not None
            or run.terminal_reason is not None
        ):
            raise QualificationRunTransitionError(
                "a new qualification run must be admitted at revision zero"
            )
        final = self._run_path(run.run_id)
        with self._lock:
            if final.exists():
                raise QualificationRunConflictError("qualification run identity already exists")
            staging = self.root / f".{run.run_id}.{uuid4().hex}.tmp"
            try:
                staging.mkdir(mode=0o700)
                for name in ("events", "commands", "artifacts"):
                    (staging / name).mkdir(mode=0o700)
                write_json_atomic(staging / "manifest.json", run.to_manifest_dict())
                _fsync_directory(staging)
                try:
                    staging.rename(final)
                except FileExistsError as exc:
                    raise QualificationRunConflictError(
                        "qualification run identity already exists"
                    ) from exc
                _fsync_directory(self.root)
            finally:
                if staging.exists():
                    shutil.rmtree(staging)
        return self.load(run.run_id)

    def load(self, run_id: str) -> QualificationRun:
        path = self._existing_run_path(run_id)
        with self._lock:
            manifest = QualificationRun.from_manifest_dict(
                _read_json(path / "manifest.json", "run manifest")
            )
            if manifest.run_id != run_id:
                raise QualificationRunIntegrityError("manifest identity does not match its path")
            events = self._read_events(path, run_id)
            artifacts = self._read_artifacts(path, run_id)
        return _replay_runtime(manifest, events, artifacts)

    def list_runs(self) -> tuple[QualificationRun, ...]:
        with self._lock:
            run_ids = tuple(
                child.name
                for child in sorted(self.root.iterdir(), key=lambda item: item.name)
                if child.is_dir() and not child.name.startswith(".")
            )
        return tuple(self.load(run_id) for run_id in run_ids)

    def transition(
        self,
        run_id: str,
        new_state: RunState,
        *,
        expected_revision: int,
        observed_at_utc: str,
        host_monotonic_ns: int,
        sim_time_ns: int | None = None,
        reason: str | None = None,
    ) -> QualificationRun:
        self._require_writable()
        with self._lock:
            current = self.load(run_id)
            if current.revision != expected_revision:
                raise QualificationRunConflictError(
                    f"run revision is {current.revision}, not {expected_revision}"
                )
            if new_state not in TRANSITIONS[current.state]:
                raise QualificationRunTransitionError(
                    f"transition {current.state.value} -> {new_state.value} is not allowed"
                )
            if new_state.terminal and (reason is None or not reason.strip()):
                raise QualificationRunTransitionError("terminal transition requires a reason")
            if not new_state.terminal and reason is not None:
                raise QualificationRunTransitionError(
                    "a nonterminal transition cannot set a terminal reason"
                )
            self._append_event_locked(
                current,
                event_type="lifecycle.state-changed",
                observed_at_utc=observed_at_utc,
                host_monotonic_ns=host_monotonic_ns,
                sim_time_ns=sim_time_ns,
                payload={
                    "from": current.state.value,
                    "to": new_state.value,
                    "reason": reason,
                },
            )
        return self.load(run_id)

    def append_event(
        self,
        run_id: str,
        *,
        event_type: str,
        observed_at_utc: str,
        host_monotonic_ns: int,
        payload: dict[str, Any],
        sim_time_ns: int | None = None,
        expected_revision: int | None = None,
    ) -> QualificationEvent:
        self._require_writable()
        if event_type == "lifecycle.state-changed":
            raise QualificationRunTransitionError("lifecycle events must use transition()")
        with self._lock:
            current = self.load(run_id)
            if current.state.terminal:
                raise QualificationRunTransitionError(
                    "terminal qualification-run history is immutable"
                )
            if expected_revision is not None and current.revision != expected_revision:
                raise QualificationRunConflictError(
                    f"run revision is {current.revision}, not {expected_revision}"
                )
            return self._append_event_locked(
                current,
                event_type=event_type,
                observed_at_utc=observed_at_utc,
                host_monotonic_ns=host_monotonic_ns,
                sim_time_ns=sim_time_ns,
                payload=payload,
            )

    def submit_command(self, command: ControlSetpoint) -> ControlSetpoint:
        self._require_writable()
        with self._lock:
            run = self.load(command.run_id)
            _validate_command_admission(command, run)
            directory = self._existing_run_path(command.run_id) / "commands"
            expected_sequence = _next_sequence(directory)
            if command.sequence != expected_sequence:
                raise QualificationRunConflictError(f"command sequence must be {expected_sequence}")
            _append_json_exclusive(
                directory / _sequence_name(command.sequence),
                command.to_dict(),
            )
        return command

    def submit_admitted_command(
        self,
        command: ControlSetpoint,
        *,
        admitted_run: QualificationRun,
    ) -> ControlSetpoint:
        """Fsync one command using the worker's single-writer admitted run snapshot."""

        self._require_writable()
        with self._lock:
            if admitted_run.run_id != command.run_id:
                raise QualificationRunConflictError("command belongs to another admitted run")
            _validate_command_admission(command, admitted_run)
            directory = self._existing_run_path(command.run_id) / "commands"
            if command.sequence < 1:
                raise QualificationRunConflictError("command sequence must be positive")
            previous = directory / _sequence_name(command.sequence - 1)
            if command.sequence > 1 and not previous.is_file():
                raise QualificationRunConflictError("command sequence would create a journal gap")
            _append_json_exclusive(
                directory / _sequence_name(command.sequence),
                command.to_dict(),
            )
        return command

    def list_commands(self, run_id: str) -> tuple[ControlSetpoint, ...]:
        path = self._existing_run_path(run_id)
        with self._lock:
            records = _read_sequence_directory(path / "commands", "command")
        commands: list[ControlSetpoint] = []
        for sequence, record in records:
            try:
                command = control_setpoint_from_dict(record)
            except SimulationContractError as exc:
                raise QualificationRunIntegrityError(str(exc)) from exc
            if command.run_id != run_id or command.sequence != sequence:
                raise QualificationRunIntegrityError(
                    "command identity/sequence does not match its journal path"
                )
            commands.append(command)
        return tuple(commands)

    def list_events(self, run_id: str) -> tuple[QualificationEvent, ...]:
        path = self._existing_run_path(run_id)
        with self._lock:
            return self._read_events(path, run_id)

    def register_artifact(
        self,
        run_id: str,
        artifact: QualificationArtifact,
    ) -> QualificationArtifact:
        self._require_writable()
        with self._lock:
            run = self.load(run_id)
            if run.state.terminal:
                raise QualificationRunTransitionError(
                    "terminal qualification-run artifacts are sealed"
                )
            directory = self._existing_run_path(run_id) / "artifacts"
            if any(existing.artifact_id == artifact.artifact_id for existing in run.artifacts):
                raise QualificationRunConflictError("artifact id is already registered")
            sequence = _next_sequence(directory)
            if artifact.sequence not in {0, sequence}:
                raise QualificationRunConflictError(f"artifact sequence must be zero or {sequence}")
            admitted = QualificationArtifact(
                artifact_id=artifact.artifact_id,
                kind=artifact.kind,
                relative_path=artifact.relative_path,
                sha256=artifact.sha256,
                byte_length=artifact.byte_length,
                source_of_record=artifact.source_of_record,
                sequence=sequence,
            )
            record = {
                "schema_version": "missioncore.qualification-artifact/v1",
                "run_id": run_id,
                **admitted.to_dict(),
            }
            _append_json_exclusive(directory / _sequence_name(sequence), record)
        return admitted

    def create_reset_episode(
        self,
        previous_run_id: str,
        *,
        run_id: str,
        episode_id: str,
        created_at_utc: str,
    ) -> QualificationRun:
        self._require_writable()
        previous = self.load(previous_run_id)
        if not previous.state.terminal:
            raise QualificationRunTransitionError(
                "reset requires the previous run to be terminal and sealed"
            )
        replacement = QualificationRun(
            run_id=run_id,
            episode_id=episode_id,
            kind=previous.kind,
            state=RunState.ADMITTED,
            scenario_generation=previous.scenario_generation,
            scenario_sha256=previous.scenario_sha256,
            profile_generation=previous.profile_generation,
            profile_sha256=previous.profile_sha256,
            mission_core_commit=previous.mission_core_commit,
            providers=previous.providers,
            host_profile_id=previous.host_profile_id,
            host_profile_sha256=previous.host_profile_sha256,
            seed=previous.seed,
            reproducibility_tier=previous.reproducibility_tier,
            authority=previous.authority,
            clock_domain=previous.clock_domain,
            created_at_utc=created_at_utc,
            parent_run_id=previous.run_id,
        )
        return self.create(replacement)

    def reconcile_interrupted(
        self,
        run_id: str,
        *,
        observed_at_utc: str,
        host_monotonic_ns: int,
    ) -> QualificationRun:
        self._require_writable()
        run = self.load(run_id)
        if run.state not in ACTIVE_RECOVERY_STATES:
            return run
        return self.transition(
            run_id,
            RunState.FAILED,
            expected_revision=run.revision,
            observed_at_utc=observed_at_utc,
            host_monotonic_ns=host_monotonic_ns,
            reason="orchestrator-recovery-interrupted",
        )

    def _append_event_locked(
        self,
        run: QualificationRun,
        *,
        event_type: str,
        observed_at_utc: str,
        host_monotonic_ns: int,
        sim_time_ns: int | None,
        payload: dict[str, Any],
    ) -> QualificationEvent:
        sequence = run.revision + 1
        event = QualificationEvent(
            run_id=run.run_id,
            sequence=sequence,
            event_type=event_type,
            observed_at_utc=observed_at_utc,
            host_monotonic_ns=host_monotonic_ns,
            sim_time_ns=sim_time_ns,
            payload=payload,
        )
        directory = self._existing_run_path(run.run_id) / "events"
        if _next_sequence(directory) != sequence:
            raise QualificationRunConflictError("event journal changed concurrently")
        _append_json_exclusive(directory / _sequence_name(sequence), event.to_dict())
        return event

    def _read_events(self, run_path: Path, run_id: str) -> tuple[QualificationEvent, ...]:
        events: list[QualificationEvent] = []
        for sequence, record in _read_sequence_directory(run_path / "events", "event"):
            try:
                event = QualificationEvent.from_dict(record)
            except SimulationContractError as exc:
                raise QualificationRunIntegrityError(str(exc)) from exc
            if event.run_id != run_id or event.sequence != sequence:
                raise QualificationRunIntegrityError(
                    "event identity/sequence does not match its journal path"
                )
            events.append(event)
        return tuple(events)

    def _read_artifacts(
        self,
        run_path: Path,
        run_id: str,
    ) -> tuple[QualificationArtifact, ...]:
        artifacts: list[QualificationArtifact] = []
        identifiers: set[str] = set()
        for sequence, record in _read_sequence_directory(
            run_path / "artifacts",
            "artifact",
        ):
            if record.get("schema_version") != "missioncore.qualification-artifact/v1":
                raise QualificationRunIntegrityError("unsupported qualification-artifact schema")
            if record.get("run_id") != run_id:
                raise QualificationRunIntegrityError("artifact run identity does not match")
            try:
                artifact = QualificationArtifact.from_dict(record)
            except SimulationContractError as exc:
                raise QualificationRunIntegrityError(str(exc)) from exc
            if artifact.sequence != sequence:
                raise QualificationRunIntegrityError(
                    "artifact sequence does not match its journal path"
                )
            if artifact.artifact_id in identifiers:
                raise QualificationRunIntegrityError("artifact journal contains a duplicate id")
            identifiers.add(artifact.artifact_id)
            artifacts.append(artifact)
        return tuple(artifacts)

    def _run_path(self, run_id: str) -> Path:
        try:
            QualificationEvent(
                run_id=run_id,
                sequence=1,
                event_type="identity.check",
                observed_at_utc="1970-01-01T00:00:00Z",
                host_monotonic_ns=0,
                sim_time_ns=None,
                payload={},
            )
        except SimulationContractError as exc:
            raise QualificationRunIntegrityError(str(exc)) from exc
        return self.root / run_id

    def _existing_run_path(self, run_id: str) -> Path:
        path = self._run_path(run_id)
        if not path.is_dir():
            raise QualificationRunNotFoundError("qualification run was not found")
        _reject_symlink(path, "qualification run")
        for name in ("events", "commands", "artifacts"):
            directory = path / name
            if not directory.is_dir():
                raise QualificationRunIntegrityError(f"run {name} journal is missing")
            _reject_symlink(directory, f"run {name} journal")
        return path

    def _require_writable(self) -> None:
        if self.read_only:
            raise QualificationRunTransitionError(
                "the qualification repository is open read-only"
            )


def _replay_runtime(
    manifest: QualificationRun,
    events: tuple[QualificationEvent, ...],
    artifacts: tuple[QualificationArtifact, ...],
) -> QualificationRun:
    state = RunState.ADMITTED
    started_at_utc: str | None = None
    ended_at_utc: str | None = None
    terminal_reason: str | None = None
    for event in events:
        if state.terminal:
            raise QualificationRunIntegrityError("event exists after terminal lifecycle state")
        if event.event_type != "lifecycle.state-changed":
            continue
        source = event.payload.get("from")
        target = event.payload.get("to")
        reason = event.payload.get("reason")
        if source != state.value or not isinstance(target, str):
            raise QualificationRunIntegrityError("lifecycle journal does not match current state")
        try:
            new_state = RunState(target)
        except ValueError as exc:
            raise QualificationRunIntegrityError("lifecycle journal has an unknown state") from exc
        if new_state not in TRANSITIONS[state]:
            raise QualificationRunIntegrityError("lifecycle journal contains an illegal transition")
        if new_state is RunState.RUNNING and started_at_utc is None:
            started_at_utc = event.observed_at_utc
        if new_state.terminal:
            if not isinstance(reason, str) or not reason.strip():
                raise QualificationRunIntegrityError("terminal lifecycle event has no reason")
            ended_at_utc = event.observed_at_utc
            terminal_reason = reason
        elif reason is not None:
            raise QualificationRunIntegrityError("nonterminal lifecycle event has a reason")
        state = new_state
    try:
        return manifest.with_runtime(
            state=state,
            revision=len(events),
            started_at_utc=started_at_utc,
            ended_at_utc=ended_at_utc,
            terminal_reason=terminal_reason,
            artifacts=artifacts,
        )
    except SimulationContractError as exc:
        raise QualificationRunIntegrityError(str(exc)) from exc


def _read_sequence_directory(
    directory: Path,
    label: str,
) -> tuple[tuple[int, dict[str, Any]], ...]:
    _reject_symlink(directory, f"{label} journal")
    records: list[tuple[int, dict[str, Any]]] = []
    expected = 1
    for path in sorted(directory.iterdir(), key=lambda item: item.name):
        _reject_regular_file(path, f"{label} record")
        if path.name != _sequence_name(expected):
            raise QualificationRunIntegrityError(f"{label} journal is not a contiguous sequence")
        records.append((expected, _read_json(path, f"{label} record")))
        expected += 1
    return tuple(records)


def _next_sequence(directory: Path) -> int:
    return len(_read_sequence_directory(directory, "journal")) + 1


def _validate_command_admission(
    command: ControlSetpoint,
    run: QualificationRun,
) -> None:
    if run.kind not in COMMANDABLE_RUN_KINDS:
        raise QualificationRunTransitionError(
            f"{run.kind.value} cannot causally accept control commands"
        )
    if run.state is not RunState.RUNNING:
        raise QualificationRunTransitionError(
            "control commands are accepted only while the run is running"
        )
    if command.authority_generation != run.authority.generation:
        raise QualificationRunConflictError("command authority generation is stale")
    ttl_ns = command.valid_until_sim_ns - command.issued_at_sim_ns
    if ttl_ns > run.authority.command_ttl_max_ns:
        raise QualificationRunTransitionError("command TTL exceeds the authority limit")


def _sequence_name(sequence: int) -> str:
    return f"{sequence:0{SEQUENCE_WIDTH}d}.json"


def _append_json_exclusive(path: Path, payload: object) -> None:
    serialized = (json.dumps(payload, ensure_ascii=False, indent=2) + "\n").encode()
    if len(serialized) > MAX_RECORD_BYTES:
        raise QualificationRunIntegrityError("journal record exceeds the size limit")
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    try:
        descriptor = os.open(path, flags, 0o600)
    except FileExistsError as exc:
        raise QualificationRunConflictError("journal sequence already exists") from exc
    try:
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(serialized)
            stream.flush()
            os.fsync(stream.fileno())
    except BaseException:
        path.unlink(missing_ok=True)
        raise
    _fsync_directory(path.parent)


def _read_json(path: Path, label: str) -> dict[str, Any]:
    _reject_regular_file(path, label)
    if path.stat().st_size > MAX_RECORD_BYTES:
        raise QualificationRunIntegrityError(f"{label} exceeds the size limit")
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise QualificationRunIntegrityError(f"{label} is not valid UTF-8 JSON") from exc
    if not isinstance(value, dict) or any(not isinstance(key, str) for key in value):
        raise QualificationRunIntegrityError(f"{label} must be a JSON object")
    return value


def _reject_regular_file(path: Path, label: str) -> None:
    try:
        mode = path.lstat().st_mode
    except FileNotFoundError as exc:
        raise QualificationRunIntegrityError(f"{label} is missing") from exc
    if not stat.S_ISREG(mode):
        raise QualificationRunIntegrityError(f"{label} must be a regular file")


def _reject_symlink(path: Path, label: str) -> None:
    try:
        mode = path.lstat().st_mode
    except FileNotFoundError as exc:
        raise QualificationRunIntegrityError(f"{label} is missing") from exc
    if stat.S_ISLNK(mode):
        raise QualificationRunIntegrityError(f"{label} must not be a symlink")


def _chmod_private(path: Path) -> None:
    with suppress(OSError):
        path.chmod(0o700)


def _fsync_directory(path: Path) -> None:
    flags = os.O_RDONLY
    if hasattr(os, "O_DIRECTORY"):
        flags |= os.O_DIRECTORY
    descriptor = os.open(path, flags)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)
