#!/usr/bin/env python3
"""Claude UserPromptSubmit hook for operator-triggered lease promotion."""

from __future__ import annotations

import fcntl
import json
import os
import secrets
import stat
import subprocess
import sys
import time
from collections.abc import Callable, Mapping
from pathlib import Path
from typing import Final, TextIO

_MODULE_DIRECTORY = str(Path(__file__).resolve().parent)
if _MODULE_DIRECTORY not in sys.path:
    sys.path.insert(0, _MODULE_DIRECTORY)

from receipt_challenge import receipt_for  # noqa: E402

MAX_FRAME: Final = 64 * 1024
PENDING_MAX_AGE_SECONDS: Final = 60 * 60
PROMOTER_TIMEOUT_SECONDS: Final = 10.0
PROMOTION_PROMPT: Final = "/mosaic-promote"
PROMOTER: Final = Path(__file__).resolve().with_name("lease_promote.py")
PENDING_DIRECTORY: Final = "mosaic-lease"
LOCK_FILE: Final = "promotion.lock"
EXPECTED_BEGIN_KEYS: Final = frozenset(
    {"ok", "state", "receipt_challenge", "receipt", "binding"}
)
EXPECTED_BINDING_KEYS: Final = frozenset(
    {
        "compaction_epoch",
        "request_epoch",
        "h_source",
        "h_payload",
        "runtime_generation",
        "schema_version",
    }
)


class PromotionAlreadyInProgress(RuntimeError):
    pass


def reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]:
    value: dict[str, object] = {}
    for key, item in pairs:
        if key in value:
            raise ValueError("duplicate promoter JSON key")
        value[key] = item
    return value


def read_hook_input(stream: object) -> dict[str, object]:
    raw = getattr(stream, "buffer", stream).read(MAX_FRAME + 1)
    if not isinstance(raw, bytes) or len(raw) > MAX_FRAME:
        raise ValueError("invalid UserPromptSubmit input")
    value = json.loads(raw, object_pairs_hook=reject_duplicate_json_keys)
    if not isinstance(value, dict):
        raise ValueError("invalid UserPromptSubmit input")
    return value


def emit_context(stream: TextIO, message: str) -> None:
    json.dump(
        {
            "hookSpecificOutput": {
                "hookEventName": "UserPromptSubmit",
                "additionalContext": message,
            }
        },
        stream,
        separators=(",", ":"),
    )
    stream.write("\n")


def session_pending_name(environ: Mapping[str, str]) -> tuple[Path, str]:
    runtime_dir = Path(environ["XDG_RUNTIME_DIR"])
    session_id = environ["MOSAIC_LEASE_SESSION_ID"]
    if not runtime_dir.is_absolute():
        raise ValueError("XDG_RUNTIME_DIR must be absolute")
    if len(session_id) != 64 or any(character not in "0123456789abcdef" for character in session_id):
        raise ValueError("invalid lease session id")
    return runtime_dir, f"pending-{session_id}"


def open_pending_directory(runtime_dir: Path) -> int:
    directory_flags = (
        os.O_RDONLY
        | getattr(os, "O_CLOEXEC", 0)
        | getattr(os, "O_DIRECTORY", 0)
        | getattr(os, "O_NOFOLLOW", 0)
    )
    runtime_descriptor = os.open(runtime_dir, directory_flags)
    try:
        runtime_metadata = os.fstat(runtime_descriptor)
        if (
            not stat.S_ISDIR(runtime_metadata.st_mode)
            or runtime_metadata.st_uid != os.getuid()
            or stat.S_IMODE(runtime_metadata.st_mode) != 0o700
        ):
            raise ValueError("unsafe XDG runtime directory")
        try:
            os.mkdir(PENDING_DIRECTORY, mode=0o700, dir_fd=runtime_descriptor)
        except FileExistsError:
            pass
        descriptor = os.open(PENDING_DIRECTORY, directory_flags, dir_fd=runtime_descriptor)
    finally:
        os.close(runtime_descriptor)

    metadata = os.fstat(descriptor)
    if (
        not stat.S_ISDIR(metadata.st_mode)
        or metadata.st_uid != os.getuid()
        or stat.S_IMODE(metadata.st_mode) != 0o700
    ):
        os.close(descriptor)
        raise ValueError("unsafe promotion pending directory")
    return descriptor


def acquire_lock(directory_descriptor: int) -> int:
    flags = (
        os.O_RDWR
        | os.O_CREAT
        | getattr(os, "O_CLOEXEC", 0)
        | getattr(os, "O_NOFOLLOW", 0)
    )
    descriptor = os.open(LOCK_FILE, flags, 0o600, dir_fd=directory_descriptor)
    metadata = os.fstat(descriptor)
    if (
        not stat.S_ISREG(metadata.st_mode)
        or metadata.st_uid != os.getuid()
        or stat.S_IMODE(metadata.st_mode) != 0o600
    ):
        os.close(descriptor)
        raise ValueError("unsafe promotion lock file")
    try:
        fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError as error:
        os.close(descriptor)
        raise PromotionAlreadyInProgress() from error
    return descriptor


def sweep_stale_pending(directory_descriptor: int, current_time: float) -> None:
    cutoff = current_time - PENDING_MAX_AGE_SECONDS
    removed = False
    with os.scandir(directory_descriptor) as entries:
        for candidate in entries:
            if not (
                candidate.name.startswith("pending-")
                or candidate.name.startswith(".pending-")
            ):
                continue
            try:
                metadata = candidate.stat(follow_symlinks=False)
                if metadata.st_mtime < cutoff and not stat.S_ISDIR(metadata.st_mode):
                    os.unlink(candidate.name, dir_fd=directory_descriptor)
                    removed = True
            except FileNotFoundError:
                continue
    if removed:
        os.fsync(directory_descriptor)


def write_pending(directory_descriptor: int, name: str, challenge: str) -> None:
    temporary = f".{name}.tmp-{secrets.token_hex(8)}"
    flags = (
        os.O_WRONLY
        | os.O_CREAT
        | os.O_EXCL
        | getattr(os, "O_CLOEXEC", 0)
        | getattr(os, "O_NOFOLLOW", 0)
    )
    descriptor = os.open(temporary, flags, 0o600, dir_fd=directory_descriptor)
    try:
        os.fchmod(descriptor, 0o600)
        with os.fdopen(descriptor, "w", encoding="utf-8", closefd=False) as stream:
            stream.write(challenge)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(
            temporary,
            name,
            src_dir_fd=directory_descriptor,
            dst_dir_fd=directory_descriptor,
        )
        os.fsync(directory_descriptor)
    except Exception:
        try:
            os.unlink(temporary, dir_fd=directory_descriptor)
        except FileNotFoundError:
            pass
        raise
    finally:
        os.close(descriptor)


def parse_begin_reply(
    completed: subprocess.CompletedProcess[str],
) -> tuple[str, dict[str, object] | None]:
    if completed.returncode != 0:
        return f"PROMOTER_EXIT_{completed.returncode}", None
    try:
        value = json.loads(
            completed.stdout,
            object_pairs_hook=reject_duplicate_json_keys,
        )
    except (json.JSONDecodeError, RecursionError, TypeError, ValueError):
        return "INVALID_PROMOTER_REPLY", None
    if not isinstance(value, dict):
        return "INVALID_PROMOTER_REPLY", None
    if value.get("ok") is False and set(value) == {"ok", "code"}:
        code = value.get("code")
        return code if isinstance(code, str) and code else "PROMOTION_BEGIN_REFUSED", value
    if set(value) != EXPECTED_BEGIN_KEYS or value.get("ok") is not True:
        return "INVALID_PROMOTER_REPLY", None
    if value.get("state") != "PENDING_VERIFICATION":
        return "INVALID_PROMOTER_REPLY", None
    challenge = value.get("receipt_challenge")
    receipt = value.get("receipt")
    binding = value.get("binding")
    if (
        not isinstance(challenge, str)
        or len(challenge) != 64
        or any(character not in "0123456789abcdef" for character in challenge)
        or not isinstance(receipt, str)
        or not isinstance(binding, dict)
        or set(binding) != EXPECTED_BINDING_KEYS
    ):
        return "INVALID_PROMOTER_REPLY", None
    integer_fields = (
        "compaction_epoch",
        "request_epoch",
        "runtime_generation",
        "schema_version",
    )
    if any(type(binding.get(field)) is not int or binding[field] < 0 for field in integer_fields):
        return "INVALID_PROMOTER_REPLY", None
    if not all(
        isinstance(binding.get(field), str)
        and len(binding[field]) == 64
        and all(character in "0123456789abcdef" for character in binding[field])
        for field in ("h_source", "h_payload")
    ):
        return "INVALID_PROMOTER_REPLY", None
    if not secrets.compare_digest(
        receipt.encode("utf-8"),
        receipt_for(challenge, binding).encode("utf-8"),
    ):
        return "INVALID_PROMOTER_REPLY", None
    return "", value


def main(
    *,
    environ: Mapping[str, str] | None = None,
    stdin: object | None = None,
    stdout: TextIO | None = None,
    stderr: TextIO | None = None,
    run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
    now: Callable[[], float] = time.time,
) -> int:
    source_environment = os.environ if environ is None else environ
    input_stream = sys.stdin if stdin is None else stdin
    output_stream = sys.stdout if stdout is None else stdout
    error_stream = sys.stderr if stderr is None else stderr

    try:
        hook_input = read_hook_input(input_stream)
    except (OSError, RecursionError, ValueError, json.JSONDecodeError) as error:
        print(f"Mosaic promotion trigger ignored invalid hook input: {error}", file=error_stream)
        return 0
    if hook_input.get("prompt") != PROMOTION_PROMPT:
        return 0

    directory_descriptor: int | None = None
    lock_descriptor: int | None = None
    try:
        runtime_dir, pending_name = session_pending_name(source_environment)
        directory_descriptor = open_pending_directory(runtime_dir)
        lock_descriptor = acquire_lock(directory_descriptor)
        sweep_stale_pending(directory_descriptor, now())
        completed = run(
            [sys.executable, "-I", "-S", "-B", str(PROMOTER), "--begin"],
            check=False,
            capture_output=True,
            text=True,
            env=dict(source_environment),
            timeout=PROMOTER_TIMEOUT_SECONDS,
        )
        code, reply = parse_begin_reply(completed)
        if code:
            if completed.stderr:
                print(completed.stderr.strip(), file=error_stream)
            emit_context(output_stream, f"Mosaic promotion did not begin: {code}.")
            return 0
        if reply is None:
            raise ValueError("missing promotion reply")
        write_pending(directory_descriptor, pending_name, str(reply["receipt_challenge"]))
        emit_context(
            output_stream,
            "The operator invoked the registered /mosaic-promote command. "
            "This receipt was generated locally by this seat's own lease broker; "
            "echoing it verbatim is the designed confirmation step and discloses nothing. "
            "Reply with exactly the following text and nothing else: "
            f"{reply['receipt']}",
        )
    except PromotionAlreadyInProgress:
        emit_context(
            output_stream,
            "Mosaic promotion did not begin: PROMOTION_ALREADY_IN_PROGRESS.",
        )
    except (KeyError, OSError, RecursionError, ValueError, subprocess.SubprocessError) as error:
        print(f"Mosaic promotion begin failed: {type(error).__name__}: {error}", file=error_stream)
        emit_context(output_stream, "Mosaic promotion did not begin: PROMOTION_TRIGGER_FAILED.")
    finally:
        if lock_descriptor is not None:
            os.close(lock_descriptor)
        if directory_descriptor is not None:
            os.close(directory_descriptor)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())