"""Durable immutable corpus generations and an atomic ACTIVE pointer."""

from __future__ import annotations

import json
import fcntl
import os
import re
import stat
import shutil
import uuid
import hashlib
import threading
from datetime import UTC, datetime
from pathlib import Path
from contextlib import contextmanager

from tht.corpus.models import CorpusManifest


_GENERATION = re.compile(r"^gen:[0-9a-f]{32}$")


class UnsafeCorpusPath(RuntimeError):
    pass


def _atomic_write(path: Path, payload: bytes) -> None:
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
    try:
        with os.fdopen(fd, "wb") as stream:
            stream.write(payload)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
        directory = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
    except BaseException:
        temporary.unlink(missing_ok=True)
        raise


class CorpusStore:
    def __init__(self, root: Path) -> None:
        self.root = Path(root)
        self.active_path = self.root / "ACTIVE"
        self._replace = os.replace
        self._fsync_directory = self._sync_root
        self._lock_state = threading.local()
        self._ensure_root()

    def _ensure_root(self) -> None:
        if self.root.is_symlink():
            raise UnsafeCorpusPath("corpus root must not be a symlink")
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        info = self.root.lstat()
        if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.getuid():
            raise UnsafeCorpusPath("corpus root is unsafe")

    @contextmanager
    def writer_lock(self):
        depth = getattr(self._lock_state, "depth", 0)
        if depth:
            self._lock_state.depth = depth + 1
            try:
                yield
            finally:
                self._lock_state.depth -= 1
            return
        lock_path = self.root / ".writer.lock"
        fd = os.open(lock_path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600)
        try:
            info = os.fstat(fd)
            if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_nlink != 1:
                raise UnsafeCorpusPath("corpus writer lock is unsafe")
            fcntl.flock(fd, fcntl.LOCK_EX)
            self._lock_state.depth = 1
            yield
        finally:
            self._lock_state.depth = 0
            fcntl.flock(fd, fcntl.LOCK_UN)
            os.close(fd)

    def generation_path(self, generation: str) -> Path:
        if not _GENERATION.fullmatch(generation):
            raise UnsafeCorpusPath("invalid corpus generation")
        path = self.root / generation.replace(":", "-")
        if path.is_symlink():
            raise UnsafeCorpusPath("generation must not be a symlink")
        return path

    def stage(
        self,
        manifest: CorpusManifest,
        materialized: dict[str, str],
        *,
        generation: str | None = None,
    ) -> str:
        generation = generation or f"gen:{uuid.uuid4().hex}"
        path = self.generation_path(generation)
        try:
            path.mkdir(mode=0o700)
        except FileExistsError:
            raise UnsafeCorpusPath("generation already exists") from None
        documents = path / "documents"
        documents.mkdir(mode=0o700)
        files: dict[str, str] = {}
        for document in manifest.documents:
            relative = f"documents/{document.document_id.removeprefix('doc:')}.md"
            _atomic_write(path / relative, materialized[document.document_id].encode("utf-8"))
            files[document.document_id] = relative
        payload = json.loads(manifest.model_dump_json())
        metadata = payload["metadata"]
        metadata["files"] = files
        payload.update({"manifest_id": generation, "metadata": metadata})
        staged = CorpusManifest.model_validate(payload)
        _atomic_write(path / "manifest.json", (staged.model_dump_json(indent=2) + "\n").encode())
        return generation

    def publish(self, generation: str) -> str:
        manifest = self.manifest(generation)
        if manifest.manifest_id != generation:
            raise UnsafeCorpusPath("manifest generation mismatch")
        if self.active_generation() == generation:
            return generation
        previous = self.active_generation()
        published_marker = self.generation_path(generation) / "PUBLISHED"
        temporary = self.active_path.with_name(f".ACTIVE.{uuid.uuid4().hex}.tmp")
        replaced = False
        try:
            _atomic_write(temporary, (generation + "\n").encode())
            self._replace(temporary, self.active_path)
            replaced = True
            self._fsync_directory()
            _atomic_write(
                published_marker,
                (datetime.now(UTC).isoformat().replace("+00:00", "Z") + "\n").encode("ascii"),
            )
        except BaseException:
            temporary.unlink(missing_ok=True)
            if replaced:
                if previous is None:
                    self.active_path.unlink(missing_ok=True)
                else:
                    rollback = self.active_path.with_name(f".ACTIVE.rollback.{uuid.uuid4().hex}.tmp")
                    _atomic_write(rollback, (previous + "\n").encode())
                    self._replace(rollback, self.active_path)
                self._sync_root()
            raise
        return generation

    def _sync_root(self) -> None:
        directory = os.open(self.root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)

    def active_generation(self) -> str | None:
        try:
            if self.active_path.is_symlink():
                raise UnsafeCorpusPath("ACTIVE must not be a symlink")
            value = self.active_path.read_text(encoding="ascii").strip()
        except FileNotFoundError:
            return None
        if not _GENERATION.fullmatch(value):
            raise UnsafeCorpusPath("ACTIVE contains an invalid generation")
        return value

    def manifest(self, generation: str) -> CorpusManifest:
        path = self.generation_path(generation)
        manifest_path = path / "manifest.json"
        if manifest_path.is_symlink():
            raise UnsafeCorpusPath("manifest must not be a symlink")
        return CorpusManifest.model_validate_json(manifest_path.read_text(encoding="utf-8"))

    def discard(self, generation: str) -> None:
        path = self.generation_path(generation)
        if path.exists():
            if path.is_symlink() or not stat.S_ISDIR(path.lstat().st_mode):
                raise UnsafeCorpusPath("generation cleanup target is unsafe")
            shutil.rmtree(path)

    def active_manifest(self) -> CorpusManifest | None:
        generation = self.active_generation()
        return self.manifest(generation) if generation else None

    def list_generations(self) -> list[str]:
        values = []
        for entry in self.root.iterdir():
            match = re.fullmatch(r"gen-([0-9a-f]{32})", entry.name)
            if match and not entry.is_symlink() and stat.S_ISDIR(entry.lstat().st_mode):
                values.append(f"gen:{match.group(1)}")
        return sorted(values, key=lambda value: self.generation_path(value).stat().st_mtime_ns)

    def published_generations(self) -> list[str]:
        active = self.active_generation()
        published = []
        for generation in self.list_generations():
            path = self.generation_path(generation)
            marker = path / "PUBLISHED"
            if generation != active and not marker.is_file():
                continue
            try:
                manifest = self.manifest(generation)
                if manifest.manifest_id != generation:
                    continue
                timestamp = marker.read_text(encoding="ascii").strip() if marker.is_file() else ""
                key = (timestamp or manifest.created_at.isoformat(), generation)
                published.append((key, generation))
            except (OSError, ValueError):
                continue
        return [generation for _, generation in sorted(published)]

    def resolve_document(self, document_id: str, generation: str | None = None) -> Path | None:
        generation = generation or self.active_generation()
        if generation is None:
            return None
        manifest = self.manifest(generation)
        relative = manifest.metadata.get("files", {}).get(document_id)
        if not isinstance(relative, str):
            return None
        parts = Path(relative).parts
        if Path(relative).is_absolute() or parts[:1] != ("documents",) or len(parts) != 2:
            raise UnsafeCorpusPath("materialized document path is unsafe")
        return self.generation_path(generation) / relative

    def read_document(self, document_id: str, generation: str | None = None) -> str | None:
        generation = generation or self.active_generation()
        if generation is None:
            return None
        manifest = self.manifest(generation)
        path = self.resolve_document(document_id, generation)
        document = next((item for item in manifest.documents if item.document_id == document_id), None)
        if path is None or document is None:
            return None
        generation_fd = os.open(self.generation_path(generation), os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
        documents_fd = fd = None
        try:
            documents_fd = os.open("documents", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=generation_fd)
            fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC, dir_fd=documents_fd)
            info = os.fstat(fd)
            if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_nlink != 1:
                raise UnsafeCorpusPath("materialized document is unsafe")
            payload = os.read(fd, info.st_size + 1)
            if len(payload) != info.st_size or "sha256:" + hashlib.sha256(payload).hexdigest() != document.content_hash:
                raise UnsafeCorpusPath("materialized document hash mismatch")
            return payload.decode("utf-8")
        except (OSError, UnicodeError) as error:
            raise UnsafeCorpusPath("materialized document read failed") from error
        finally:
            if fd is not None:
                os.close(fd)
            if documents_fd is not None:
                os.close(documents_fd)
            os.close(generation_fd)

    def materialize_document(
        self, document_id: str, destination: Path, generation: str | None = None,
    ) -> Path | None:
        content = self.read_document(document_id, generation)
        if content is None:
            return None
        destination = Path(destination)
        destination.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
        _atomic_write(destination, content.encode("utf-8"))
        destination.chmod(0o400)
        return destination