from __future__ import annotations

import fcntl
import threading
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from pathlib import Path
from typing import cast

from pydantic import BaseModel, ConfigDict

from checkin_cli.nutrition_onboarding_fs import (
    atomic_write_private_json,
    create_private_file,
    ensure_private_profile_directory,
    fsync_directory,
    read_private_json,
    validate_private_directory,
    validate_private_file,
    validate_profile_path,
)


class StoreConflictError(ValueError):
    pass


class PublicationRecord(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    session_id: str
    generation: int
    state: str
    payload: dict[str, object]
    message_id: int | None = None


class NutritionOnboardingStore:
    def __init__(self, *, root: Path, profile_root: Path | None = None) -> None:
        self.root = Path(root)
        self.profile_root = None if profile_root is None else Path(profile_root)
        self.lock_path = self.root / ".lock"
        self.state_path = self.root / "session.json"
        self.transient_path = self.root / "transient"
        self._thread_lock = threading.RLock()

    def _validate_root(self) -> None:
        if self.profile_root is not None:
            validate_profile_path(self.root, self.profile_root)

    def initialize(self) -> None:
        self._validate_root()
        if self.root.is_symlink():
            raise ValueError("store root must not be a symlink")
        if self.profile_root is not None:
            ensure_private_profile_directory(self.root, self.profile_root)
        else:
            self.root.mkdir(mode=0o700, parents=True, exist_ok=True)
            self.root.chmod(0o700)
        self._validate_root()
        validate_private_directory(self.root)
        create_private_file(self.lock_path, b"")
        if not self.state_path.exists():
            self._atomic_write({"schema_version": "1.0", "sessions": {}})
        else:
            validate_private_file(self.state_path)
        if self.profile_root is not None:
            ensure_private_profile_directory(self.transient_path, self.profile_root)
        else:
            self.transient_path.mkdir(mode=0o700, exist_ok=True)
            self.transient_path.chmod(0o700)
        validate_private_directory(self.transient_path)

    def reload(self) -> "NutritionOnboardingStore":
        self._validate_root()
        validate_private_directory(self.root)
        validate_private_file(self.lock_path)
        validate_private_file(self.state_path)
        self._read_state()
        return self

    @contextmanager
    def _locked(self) -> Iterator[None]:
        self.initialize()
        with self._thread_lock:
            with self.lock_path.open("r+b") as handle:
                fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
                try:
                    yield
                finally:
                    fcntl.flock(handle.fileno(), fcntl.LOCK_UN)

    @contextmanager
    def locked(self) -> Iterator[None]:
        with self._locked():
            yield

    def compare_and_swap_generation(
        self,
        *,
        session_id: str,
        expected_generation: int,
        next_generation: int,
    ) -> PublicationRecord:
        if next_generation != expected_generation + 1:
            raise ValueError("next generation must increment by one")
        with self._locked():
            document = self._read_state()
            sessions = self._sessions(document)
            current = sessions.get(session_id)
            generation = 0 if current is None else self._record(current).generation
            if generation != expected_generation:
                raise StoreConflictError("stale onboarding generation")
            record = PublicationRecord(
                session_id=session_id,
                generation=next_generation,
                state="ACTIVE",
                payload={},
            )
            sessions[session_id] = record.model_dump(mode="json")
            self._atomic_write(document)
            return record

    def mark_prepared(
        self,
        *,
        session_id: str,
        generation: int,
        payload: Mapping[str, object],
    ) -> PublicationRecord:
        with self._locked():
            document = self._read_state()
            sessions = self._sessions(document)
            current = sessions.get(session_id)
            if current is not None:
                existing = self._record(current)
                if existing.generation > generation:
                    raise StoreConflictError("stale onboarding generation")
                if (
                    existing.generation == generation
                    and existing.state in {"COMMITTED", "UNCERTAIN"}
                ):
                    return existing
                if (
                    existing.generation == generation
                    and existing.state == "PREPARED"
                ):
                    raise StoreConflictError(
                        "onboarding publication generation is already prepared"
                    )
            record = PublicationRecord(
                session_id=session_id,
                generation=generation,
                state="PREPARED",
                payload=dict(payload),
            )
            sessions[session_id] = record.model_dump(mode="json")
            self._atomic_write(document)
            return record

    prepare_publication = mark_prepared

    def mark_committed(
        self,
        *,
        session_id: str,
        generation: int,
        message_id: int | None = None,
    ) -> PublicationRecord:
        return self._transition(
            session_id=session_id,
            generation=generation,
            expected_state="PREPARED",
            next_state="COMMITTED",
            message_id=message_id,
        )

    commit_publication = mark_committed

    def mark_uncertain(
        self,
        *,
        session_id: str,
        generation: int,
    ) -> PublicationRecord:
        return self._transition(
            session_id=session_id,
            generation=generation,
            expected_state="PREPARED",
            next_state="UNCERTAIN",
        )

    mark_publication_uncertain = mark_uncertain

    def load_session(self, session_id: str) -> PublicationRecord:
        with self._locked():
            current = self._sessions(self._read_state()).get(session_id)
            if current is None:
                raise ValueError(f"unknown onboarding session: {session_id}")
            return self._record(current)

    def deliver(
        self,
        *,
        session_id: str,
        generation: int,
        provider: Callable[[dict[str, object]], object],
    ) -> PublicationRecord:
        prepared = self.mark_prepared(
            session_id=session_id,
            generation=generation,
            payload={"session_id": session_id, "generation": generation},
        )
        if prepared.state in {"COMMITTED", "UNCERTAIN"}:
            return prepared
        try:
            result = provider(dict(prepared.payload))
        except Exception:  # provider exception means side effects are unknown
            return self.mark_uncertain(session_id=session_id, generation=generation)
        message_id = result if isinstance(result, int) and not isinstance(result, bool) else None
        return self.mark_committed(
            session_id=session_id,
            generation=generation,
            message_id=message_id,
        )

    def purge_transient(self) -> int:
        self.initialize()
        removed = 0
        for path in self.transient_path.iterdir():
            validate_private_file(path)
            path.unlink()
            removed += 1
        fsync_directory(self.transient_path)
        return removed

    def _transition(
        self,
        *,
        session_id: str,
        generation: int,
        expected_state: str,
        next_state: str,
        message_id: int | None = None,
    ) -> PublicationRecord:
        with self._locked():
            document = self._read_state()
            sessions = self._sessions(document)
            current = sessions.get(session_id)
            if current is None:
                raise ValueError(f"unknown onboarding session: {session_id}")
            existing = self._record(current)
            if existing.generation != generation or existing.state != expected_state:
                raise StoreConflictError("stale publication transition")
            record = existing.model_copy(
                update={"state": next_state, "message_id": message_id},
            )
            sessions[session_id] = record.model_dump(mode="json")
            self._atomic_write(document)
            return record

    @staticmethod
    def _sessions(document: dict[str, object]) -> dict[str, object]:
        sessions = document.get("sessions")
        if not isinstance(sessions, dict):
            raise ValueError("invalid onboarding store")
        return cast(dict[str, object], sessions)

    @staticmethod
    def _record(value: object) -> PublicationRecord:
        return PublicationRecord.model_validate(value)

    def _read_state(self) -> dict[str, object]:
        value = read_private_json(self.state_path)
        if not isinstance(value, dict) or value.get("schema_version") != "1.0":
            raise ValueError("invalid onboarding store")
        self._sessions(value)
        return value

    def _atomic_write(self, document: dict[str, object]) -> None:
        atomic_write_private_json(self.state_path, document)