#!/usr/bin/env python3
"""Atomically persist the Python profile operation journal."""

from __future__ import annotations

import argparse
import fcntl
import hashlib
import json
import os
import tempfile
from pathlib import Path


class JournalError(ValueError):
    pass


TRANSITIONS = {
    None: {"profile_selected"},
    "profile_selected": {"checkpoint_ready", "rolled_back"},
    "checkpoint_ready": {"installing", "rolled_back"},
    "installing": {"installed", "rollback_required"},
    "rollback_required": {"rolled_back"},
    "installed": set(),
    "rolled_back": set(),
}


def canonical(record: dict) -> bytes:
    return (json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n").encode()


def digest_record(record: dict) -> str:
    unsigned = {key: value for key, value in record.items() if key != "record_sha256"}
    return hashlib.sha256(canonical(unsigned)).hexdigest()


def read_verified(path: Path) -> dict | None:
    if not path.exists():
        return None
    record = json.loads(path.read_text(encoding="utf-8"))
    if record.get("record_sha256") != digest_record(record):
        raise JournalError("JOURNAL_DIGEST_INVALID")
    return record


def next_record(current: dict | None, expected_generation: int, operation_id: str,
                profile: str, manifest_sha256: str, state: str,
                checkpoint_id: str | None) -> dict:
    generation = current.get("generation", 0) if current else 0
    if generation != expected_generation:
        raise JournalError("JOURNAL_GENERATION_STALE")
    prior_state = current.get("state") if current else None
    if state not in TRANSITIONS.get(prior_state, set()):
        raise JournalError("JOURNAL_TRANSITION_INVALID")
    if current and (current["operation_id"] != operation_id or current["selected_profile"] != profile
                    or current["manifest_sha256"] != manifest_sha256):
        raise JournalError("JOURNAL_IDENTITY_CHANGED")
    inherited_checkpoint = current.get("checkpoint_id") if current else None
    checkpoint = checkpoint_id or inherited_checkpoint
    if state in {"checkpoint_ready", "installing", "installed", "rollback_required", "rolled_back"} \
            and prior_state != "profile_selected" and not checkpoint:
        raise JournalError("CHECKPOINT_ID_REQUIRED")
    if state in {"checkpoint_ready", "installing", "installed", "rollback_required"} and not checkpoint:
        raise JournalError("CHECKPOINT_ID_REQUIRED")
    record = {
        "schema": "com.cezen.nexus-one.python-profile-journal.v1",
        "generation": generation + 1,
        "operation_id": operation_id,
        "selected_profile": profile,
        "manifest_sha256": manifest_sha256,
        "state": state,
        "checkpoint_id": checkpoint,
        "previous_record_sha256": current.get("record_sha256") if current else None,
    }
    record["record_sha256"] = digest_record(record)
    return record


def atomic_write(path: Path, record: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    try:
        os.fchmod(fd, 0o600)
        with os.fdopen(fd, "wb", closefd=True) as stream:
            stream.write(canonical(record))
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
        directory_fd = os.open(path.parent, os.O_RDONLY)
        try:
            os.fsync(directory_fd)
        finally:
            os.close(directory_fd)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def update(path: Path, **values) -> dict:
    lock_path = path.with_suffix(path.suffix + ".lock")
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    with lock_path.open("a+b") as lock:
        os.chmod(lock_path, 0o600)
        fcntl.flock(lock, fcntl.LOCK_EX)
        current = read_verified(path)
        record = next_record(current, **values)
        atomic_write(path, record)
        return record


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=Path)
    parser.add_argument("--expected-generation", type=int, required=True)
    parser.add_argument("--operation-id", required=True)
    parser.add_argument("--profile", choices=("cpu-only", "nvidia-cu128"), required=True)
    parser.add_argument("--manifest-sha256", required=True)
    parser.add_argument("--state", choices=tuple(state for state in TRANSITIONS if state), required=True)
    parser.add_argument("--checkpoint-id")
    args = parser.parse_args()
    try:
        record = update(
            args.path, expected_generation=args.expected_generation,
            operation_id=args.operation_id, profile=args.profile,
            manifest_sha256=args.manifest_sha256, state=args.state,
            checkpoint_id=args.checkpoint_id,
        )
    except (JournalError, json.JSONDecodeError) as error:
        print(json.dumps({"written": False, "code": str(error)}, sort_keys=True))
        return 1
    print(json.dumps({"written": True, **record}, sort_keys=True))
    return 0


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