#!/usr/bin/env python3
from __future__ import annotations

import fcntl
import hashlib
import json
import os
import shutil
import stat
import tempfile
from dataclasses import dataclass
from pathlib import Path

SOURCE_OPENED_HOOK = lambda _name, _fd: None

FILES = {
    "overdeck-seat-scope-entry": 0o755,
    "seat_scope_entry.py": 0o644,
    "seat_common.py": 0o644,
    "seat_implementer_identity.py": 0o644,
    "overdeck-seat-implementer-exec": 0o755,
    "seat_implementer_exec.py": 0o644,
    "overdeck-seat-tmux-mediator": 0o755,
    "seat_tmux_mediator.py": 0o644,
    "seat_execution_broker.py": 0o644,
    "seat_execution_client.py": 0o644,
    "overdeck-seat-execution-client": 0o755,
    "agent-session-reap-close.py": 0o755,
    "seat-launcher": 0o755,
}
STABLE = {
    "scope": ("usr/local/bin/overdeck-seat-scope-entry", "overdeck-seat-scope-entry", 0o755),
    "implementer": ("usr/local/bin/overdeck-seat-implementer-exec", "overdeck-seat-implementer-exec", 0o755),
    "mediator": ("usr/local/bin/overdeck-seat-tmux-mediator", "overdeck-seat-tmux-mediator", 0o755),
    "client": ("usr/local/bin/overdeck-seat-execution-client", "overdeck-seat-execution-client", 0o755),
    "sudoers": ("etc/sudoers.d/overdeck-seat-scope-entry", "sudoers", 0o440),
}


class InstallError(RuntimeError):
    pass


class InstallInterrupted(BaseException):
    pass


class RollbackIncomplete(InstallError):
    pass


@dataclass(frozen=True)
class InstallResult:
    version: str
    version_root: str


def trusted_owner(st: os.stat_result) -> bool:
    return st.st_uid in (0, os.getuid())


def _root(root: str, relative: str) -> Path:
    return Path(root) / relative


def _open_sources(source_dir: str, invoking_uid: int) -> dict[str, tuple[int, os.stat_result]]:
    directory = os.open(source_dir, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
    opened = {}
    try:
        for name, mode in FILES.items():
            try:
                fd = os.open(name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory)
            except OSError as exc:
                raise InstallError(f"source-invalid:{name}:{exc.errno}") from None
            st = os.fstat(fd)
            if (not stat.S_ISREG(st.st_mode) or st.st_uid != invoking_uid or st.st_nlink != 1
                    or st.st_size > 16 * 1024 * 1024 or stat.S_IMODE(st.st_mode) != mode):
                os.close(fd)
                raise InstallError(f"source-invalid:{name}")
            SOURCE_OPENED_HOOK(name, fd)
            opened[name] = (fd, st)
        return opened
    except BaseException:
        for fd, _ in opened.values():
            os.close(fd)
        raise
    finally:
        os.close(directory)


def _fd_bytes(fd: int, expected: os.stat_result) -> bytes:
    os.lseek(fd, 0, os.SEEK_SET)
    chunks = []
    while True:
        chunk = os.read(fd, 131072)
        if not chunk:
            break
        chunks.append(chunk)
    after = os.fstat(fd)
    identity = lambda st: (st.st_dev, st.st_ino, st.st_mode, st.st_uid, st.st_gid, st.st_size,
                           st.st_mtime_ns, st.st_ctime_ns, st.st_nlink)
    if identity(after) != identity(expected):
        raise InstallError("source-changed")
    return b"".join(chunks)


def _snapshot(source_dir: str, invoking_uid: int) -> tuple[dict[str, bytes], str]:
    opened = _open_sources(source_dir, invoking_uid)
    try:
        data = {name: _fd_bytes(fd, st) for name, (fd, st) in opened.items()}
    finally:
        for fd, _ in opened.values():
            os.close(fd)
    digest = hashlib.sha256("\n".join(hashlib.sha256(data[n]).hexdigest() for n in FILES).encode() + b"\n").hexdigest()
    return data, digest


def source_digest(source_dir: str, invoking_uid: int) -> str:
    return _snapshot(source_dir, invoking_uid)[1]


def acquire_lock(root: str) -> int:
    state = _root(root, "var/lib/overdeck")
    state.mkdir(parents=True, exist_ok=True)
    fd = os.open(state / "seat-install.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        os.close(fd)
        raise InstallError("install-busy")
    return fd


def _write_json_atomic(path: Path, value: dict) -> None:
    tmp = path.with_name(path.name + ".new")
    with open(tmp, "x", encoding="utf-8") as fh:
        json.dump(value, fh, sort_keys=True)
        fh.write("\n")
        fh.flush()
        os.fsync(fh.fileno())
    os.replace(tmp, path)


def _validate_stable(path: Path) -> None:
    try:
        st = path.lstat()
    except FileNotFoundError:
        return
    if not stat.S_ISREG(st.st_mode) or not trusted_owner(st):
        raise InstallError(f"stable-invalid:{path}")


def restore_item(item: dict) -> None:
    dst = Path(item["dst"])
    if dst.is_symlink() or dst.exists():
        if dst.is_dir() and not dst.is_symlink():
            shutil.rmtree(dst)
        else:
            dst.unlink()
    backup = Path(item["backup"])
    if item["existed"]:
        shutil.copy2(backup, dst, follow_symlinks=False)


def _rollback(journal_path: Path, journal: dict) -> None:
    failures = []
    for item in reversed(journal["touched"]):
        try:
            restore_item(item)
        except Exception as exc:
            failures.append(f"{item['name']}:{exc}")
    active = Path(journal["active"])
    try:
        if journal["active_existed"]:
            shutil.copy2(journal["active_backup"], active)
        elif active.exists() or active.is_symlink():
            active.unlink()
    except Exception as exc:
        failures.append(f"activation:{exc}")
    if failures:
        journal["rollbackFailures"] = failures
        _write_json_atomic(journal_path, journal)
        raise RollbackIncomplete("rollback-incomplete:" + ",".join(failures))
    shutil.rmtree(journal["transaction_dir"], ignore_errors=True)
    journal_path.unlink(missing_ok=True)


def _recover(root: str) -> None:
    path = _root(root, "var/lib/overdeck/seat-install-transaction.json")
    if not path.exists():
        return
    with open(path, encoding="utf-8") as fh:
        journal = json.load(fh)
    _rollback(path, journal)


def _validate_existing_version(target: Path, data: dict[str, bytes]) -> None:
    st = target.lstat()
    if not stat.S_ISDIR(st.st_mode) or not trusted_owner(st) or stat.S_IMODE(st.st_mode) & 0o022:
        raise InstallError("version-invalid")
    expected = set(FILES) | {"sudoers", "seat-scope-entry.manifest.json"}
    if {p.name for p in target.iterdir()} != expected:
        raise InstallError("version-mismatch")
    for name, content in data.items():
        p = target / name
        st = p.lstat()
        if not stat.S_ISREG(st.st_mode) or st.st_nlink != 1 or p.read_bytes() != content:
            raise InstallError("version-mismatch")


def install(*, source_dir: str, root: str = "/", invoking_uid: int, invoking_user: str,
            fault: str = "") -> InstallResult:
    lock = acquire_lock(root)
    try:
        _recover(root)
        data, version = _snapshot(source_dir, invoking_uid)
        base = _root(root, "usr/local/lib/overdeck")
        versions = base / "versions"
        state = _root(root, "var/lib/overdeck")
        for path, mode in ((base, 0o755), (versions, 0o755), (state, 0o755),
                           (_root(root, "usr/local/bin"), 0o755), (_root(root, "etc/sudoers.d"), 0o755)):
            path.mkdir(parents=True, exist_ok=True)
            path.chmod(mode)
        txdir = Path(tempfile.mkdtemp(prefix="seat-install-", dir=state))
        txdir.chmod(0o700)
        stage = txdir / "stage"
        stage.mkdir(mode=0o700)
        for name, content in data.items():
            fd = os.open(stage / name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, FILES[name])
            with os.fdopen(fd, "wb") as fh:
                fh.write(content); fh.flush(); os.fsync(fh.fileno())
        rules = (f"{invoking_user} ALL=(root) NOPASSWD: /usr/local/bin/overdeck-seat-scope-entry *\n"
                 f"{invoking_user} ALL=(root) NOPASSWD: /usr/local/bin/overdeck-seat-tmux-mediator *\n"
                 f"{invoking_user} ALL=(root) NOPASSWD: /usr/local/bin/overdeck-seat-implementer-exec *\n").encode()
        (stage / "sudoers").write_bytes(rules); (stage / "sudoers").chmod(0o440)
        manifest = {"ssh_user": invoking_user, **{f"{name}_sha256": hashlib.sha256(content).hexdigest() for name, content in data.items()}}
        (stage / "seat-scope-entry.manifest.json").write_text(json.dumps(manifest, sort_keys=True) + "\n")
        (stage / "seat-scope-entry.manifest.json").chmod(0o600)
        target = versions / version
        if target.exists() or target.is_symlink():
            _validate_existing_version(target, data)
            shutil.rmtree(stage)
        else:
            os.rename(stage, target)
            target.chmod(0o755)
        active = base / "seat-authority-active"
        journal_path = state / "seat-install-transaction.json"
        active_backup = txdir / "active"
        active_existed = active.exists()
        if active_existed:
            shutil.copy2(active, active_backup)
        journal = {"transaction_dir": str(txdir), "touched": [], "active": str(active),
                   "active_existed": active_existed, "active_backup": str(active_backup)}
        _write_json_atomic(journal_path, journal)
        try:
            for name, (relative, source_name, mode) in STABLE.items():
                dst = _root(root, relative); _validate_stable(dst)
                backup = txdir / f"backup-{name}"
                existed = dst.exists()
                if existed: shutil.copy2(dst, backup, follow_symlinks=False)
                item = {"name": name, "dst": str(dst), "backup": str(backup), "existed": existed}
                journal["touched"].append(item); _write_json_atomic(journal_path, journal)
                tmp = dst.with_name(dst.name + f".new-{os.getpid()}")
                shutil.copyfile(target / source_name, tmp, follow_symlinks=False); tmp.chmod(mode); os.replace(tmp, dst)
                if fault == f"interrupt-after-{name}": raise InstallInterrupted()
                if fault == f"after-{name}": raise InstallError(fault)
            active_tmp = txdir / "active-new"; active_tmp.write_text(version + "\n"); active_tmp.chmod(0o600)
            if fault == "after-activation": raise InstallError(fault)
            os.replace(active_tmp, active)
            journal_path.unlink(); shutil.rmtree(txdir)
            return InstallResult(version, str(target))
        except InstallInterrupted:
            raise
        except BaseException:
            _rollback(journal_path, journal)
            raise
    finally:
        os.close(lock)