#!/usr/bin/env python3
"""kb enforce installer: per-agent providers for the data-olympus enforcement gate.

Each provider idempotently installs/removes MARKER-tagged enforcement wiring for
one coding agent, backs up before editing, and reports status/doctor. Subcommands:
install | uninstall | status | doctor. Select an agent with --agent.
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
import time
import urllib.request
from pathlib import Path

MARKER = "data-olympus-enforce"
SHIM_VERSION = "1"
HOOK_BIN = str(Path(__file__).resolve().parent / "kb-enforce-hook")
PLUGIN_SRC = Path(__file__).resolve().parent / "opencode" / "data-olympus-gate.ts"
PLUGIN_NAME = "data-olympus-gate.ts"


def _home() -> str:
    override = os.getenv("KB_ENFORCE_HOME")
    return override if override else os.path.expanduser("~")


def _backup(target: Path) -> None:
    if target.exists():
        ts = time.strftime("%Y%m%d-%H%M%S")
        shutil.copy2(target, target.with_suffix(target.suffix + f".kb-bak-{ts}"))


def _load_json(target: Path) -> dict:
    if target.exists() and target.read_text().strip():
        return json.loads(target.read_text())
    return {}


def _doctor_endpoint() -> tuple[bool, str]:
    endpoint = os.getenv("KB_ENDPOINT", "http://localhost:8080")
    try:
        with urllib.request.urlopen(f"{endpoint}/api/v1/health", timeout=5) as r:
            ok = r.status == 200
    except Exception as exc:  # noqa: BLE001 - report any failure
        return False, f"cannot reach {endpoint}: {exc}"
    return ok, f"endpoint {endpoint} reachable={ok}"


def _hook_bin_in_worktree(hook_bin: str) -> bool:
    """True when the dispatcher path resolves inside a git worktree checkout
    (a `.worktrees/` or `.claude/worktrees/` segment). An install performed from
    such a checkout dangles after the worktree is pruned and then silently fails
    open, so doctor warns about it."""
    parts = Path(hook_bin).resolve().parts
    if ".worktrees" in parts:
        return True
    # `.claude/worktrees/<...>`: a `.claude` segment immediately followed by
    # `worktrees`.
    return any(
        parts[i] == ".claude" and parts[i + 1] == "worktrees"
        for i in range(len(parts) - 1)
    )


def _doctor_hook_bin() -> tuple[bool, list[str]]:
    """Verify the installed hook dispatcher exists and is executable, and warn
    when it resolves inside a worktree. Returns (ok, messages)."""
    msgs: list[str] = []
    exists = os.path.isfile(HOOK_BIN)
    executable = exists and os.access(HOOK_BIN, os.X_OK)
    if not exists:
        msgs.append(f"hook command MISSING at {HOOK_BIN}")
    elif not executable:
        msgs.append(f"hook command present but NOT executable at {HOOK_BIN}")
    else:
        msgs.append(f"hook command present and executable at {HOOK_BIN}")
    if _hook_bin_in_worktree(HOOK_BIN):
        msgs.append(
            f"WARNING: hook command resolves inside a worktree ({HOOK_BIN}); this "
            "install will dangle and silently fail open once the worktree is "
            "pruned. Re-run `kb enforce install` from the main checkout."
        )
    return executable, msgs


def _managed_versions_in_hooks(data: dict) -> set:
    """The set of managed-marker versions present in a JSON hooks map."""
    return {
        h[MARKER]
        for blocks in data.get("hooks", {}).values()
        for block in blocks
        for h in block.get("hooks", [])
        if MARKER in h
    }


class HookFileProvider:
    """A provider that writes MARKER-tagged hook entries into a JSON hooks map
    inside a target file (the map lives under the top-level 'hooks' key).

    events: list of (event_name, dispatcher_mode, matcher_or_None).
    dialect: passed to kb-enforce-hook as '--dialect <dialect>'; omitted when
    'claude' so Claude's slice-1 command form ('<hook> <mode>') is preserved.
    """

    tier = "hard"

    def __init__(self, name: str, default_target: Path, events: list,
                 dialect: str = "claude", note: str = "") -> None:
        self.name = name
        self._default_target = default_target
        self._events = events
        self._dialect = dialect
        self._note = note
        self.enforce_mode = "hard"

    def default_target(self) -> Path:
        return self._default_target

    def _command(self, mode: str) -> str:
        # Always thread --agent so the consult audit records the correct
        # per-agent identity. --dialect is still suppressed for claude (default)
        # to keep its slice-1 command form.
        dialect = "" if self._dialect == "claude" else f" --dialect {self._dialect}"
        return f"{HOOK_BIN} {mode}{dialect} --agent {self.name}"

    def _managed_block(self, mode: str, matcher: str | None) -> dict:
        entry = {"type": "command", "command": self._command(mode), MARKER: SHIM_VERSION}
        block: dict = {"hooks": [entry]}
        if matcher is not None:
            block["matcher"] = matcher
        return block

    @staticmethod
    def _strip_managed(hooks: dict) -> dict:
        out: dict = {}
        for event, blocks in hooks.items():
            kept = []
            for block in blocks:
                kh = [h for h in block.get("hooks", []) if MARKER not in h]
                if kh:
                    nb = dict(block)
                    nb["hooks"] = kh
                    kept.append(nb)
            if kept:
                out[event] = kept
        return out

    def install(self, target: Path) -> int:
        data = _load_json(target)
        _backup(target)
        hooks = self._strip_managed(data.get("hooks", {}))
        events = self._events
        if self.enforce_mode == "soft":
            events = [(e, m, mt) for (e, m, mt) in self._events if m != "pre-tool"]
        for event, mode, matcher in events:
            hooks.setdefault(event, []).append(self._managed_block(mode, matcher))
        data["hooks"] = hooks
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(json.dumps(data, indent=2) + "\n")
        print(
            f"installed data-olympus enforcement (v{SHIM_VERSION}) into {target} "
            f"[{self.name}, tier={self.tier}]"
        )
        if self._note:
            print(self._note)
        return 0

    def uninstall(self, target: Path) -> int:
        data = _load_json(target)
        if "hooks" not in data:
            print("nothing to uninstall")
            return 0
        _backup(target)
        data["hooks"] = self._strip_managed(data["hooks"])
        if not data["hooks"]:
            del data["hooks"]
        target.write_text(json.dumps(data, indent=2) + "\n")
        print(f"uninstalled data-olympus enforcement from {target} [{self.name}]")
        return 0

    def status(self, target: Path) -> int:
        data = _load_json(target)
        versions = _managed_versions_in_hooks(data)
        if not versions:
            print(f"{self.name}: not installed")
            return 0
        stale = " (stale; run `kb enforce install`)" if SHIM_VERSION not in versions else ""
        print(f"{self.name}: installed, tier={self.tier}, versions={sorted(versions)}{stale}")
        return 0

    def doctor(self, target: Path) -> int:
        """Doctor now checks three things beyond endpoint reachability:
        the managed marker/version is present in the LIVE settings file, the hook
        dispatcher exists and is executable, and the dispatcher is not installed
        from a worktree (which dangles after pruning and fails open)."""
        ok_ep, msg_ep = _doctor_endpoint()
        print(f"doctor [{self.name}]: {msg_ep}")

        data = _load_json(target)
        versions = _managed_versions_in_hooks(data)
        if not versions:
            print(f"doctor [{self.name}]: managed hook NOT installed in {target}")
            ok_marker = False
        elif SHIM_VERSION not in versions:
            print(f"doctor [{self.name}]: managed hook in {target} is STALE "
                  f"(found {sorted(versions)}, want {SHIM_VERSION}); run "
                  "`kb enforce install`")
            ok_marker = False
        else:
            print(f"doctor [{self.name}]: managed hook v{SHIM_VERSION} present in {target}")
            ok_marker = True

        ok_bin, bin_msgs = _doctor_hook_bin()
        for m in bin_msgs:
            print(f"doctor [{self.name}]: {m}")

        # A worktree-installed hook is a real problem even if the file is currently
        # present and executable, so fail doctor on it.
        ok_worktree = not _hook_bin_in_worktree(HOOK_BIN)
        return 0 if (ok_ep and ok_marker and ok_bin and ok_worktree) else 1


def _claude_provider() -> HookFileProvider:
    return HookFileProvider(
        name="claude-code",
        default_target=Path(_home()) / ".claude" / "settings.json",
        events=[
            ("SessionStart", "session-start", None),
            ("UserPromptSubmit", "user-prompt", None),
            # Anchored: the matcher is a regex, so an un-anchored alternation
            # substring-matches unrelated tools (e.g. "Bash" matches "BashOutput",
            # "Edit" matches "NotebookEditOther"). ^(...)$ gates exactly these.
            ("PreToolUse", "pre-tool", "^(Edit|Write|MultiEdit|NotebookEdit|Bash)$"),
            ("Stop", "stop", None),
        ],
        dialect="claude",
    )


CODEX_TRUST_NOTE = (
    "NOTE (codex): Codex requires this hook to be trusted before it runs. On the "
    "next `codex` start you will be prompted to trust it, or run codex with "
    "`--dangerously-bypass-hook-trust` for vetted automation. The trust hash is "
    "persisted under [hooks.state] in ~/.codex/config.toml."
)


def _codex_provider() -> HookFileProvider:
    return HookFileProvider(
        name="codex",
        default_target=Path(_home()) / ".codex" / "hooks.json",
        events=[
            ("SessionStart", "session-start", None),
            ("UserPromptSubmit", "user-prompt", None),
            # Anchored (see the claude provider): keep the alternation exact so
            # "Bash" does not also gate "BashOutput".
            ("PreToolUse", "pre-tool", "^(Edit|Write|MultiEdit|Bash)$"),
            ("Stop", "stop", None),
        ],
        dialect="claude",  # Codex shares Claude's exit-2 deny contract
        note=CODEX_TRUST_NOTE,
    )


def _gemini_provider() -> HookFileProvider:
    return HookFileProvider(
        name="gemini",
        default_target=Path(_home()) / ".gemini" / "settings.json",
        events=[
            ("SessionStart", "session-start", None),
            ("BeforeAgent", "user-prompt", None),  # BeforeAgent carries `prompt`
            # Anchored to avoid substring matches against other Gemini tools.
            ("BeforeTool", "pre-tool", "^(write_file|replace|run_shell_command)$"),
            ("Stop", "stop", None),
        ],
        dialect="gemini",
    )


class OpenCodeProvider:
    """OpenCode provider. Unlike the hard shell-hook agents, OpenCode is wired by
    dropping a managed TypeScript plugin file into the plugin directory. Install
    copies the bundled template; uninstall removes it only if it still carries the
    managed marker, so operator-authored plugins in the same dir are untouched.
    """

    name = "opencode"
    tier = "hard"

    def default_target(self) -> Path:
        return Path(_home()) / ".config" / "opencode" / "plugin"

    def install(self, target: Path) -> int:
        target.mkdir(parents=True, exist_ok=True)
        dest = target / PLUGIN_NAME
        if dest.exists():
            _backup(dest)
        shutil.copy2(PLUGIN_SRC, dest)
        print(
            f"installed data-olympus enforcement (v{SHIM_VERSION}) into {dest} "
            f"[opencode, tier=hard]"
        )
        return 0

    def uninstall(self, target: Path) -> int:
        dest = target / PLUGIN_NAME
        if dest.exists() and "data-olympus-enforce (managed)" in dest.read_text():
            dest.unlink()
            print(f"uninstalled data-olympus enforcement from {dest} [opencode]")
        else:
            print("nothing to uninstall")
        return 0

    def status(self, target: Path) -> int:
        import re
        dest = target / PLUGIN_NAME
        if not (dest.exists() and "data-olympus-enforce (managed)" in dest.read_text()):
            print("opencode: not installed")
            return 0
        m = re.search(r"data-olympus-enforce \(managed\) v(\d+)", dest.read_text())
        ver = m.group(1) if m else "?"
        stale = " (stale; run `kb enforce install`)" if ver != SHIM_VERSION else ""
        print(f"opencode: installed, tier=hard, versions=['{ver}']{stale}")
        return 0

    def doctor(self, target: Path) -> int:
        ok_ep, msg = _doctor_endpoint()
        print(f"doctor [opencode]: {msg}")
        dest = target / PLUGIN_NAME
        installed = dest.exists() and "data-olympus-enforce (managed)" in dest.read_text()
        if installed:
            print(f"doctor [opencode]: managed plugin present at {dest}")
        else:
            print(f"doctor [opencode]: managed plugin NOT installed at {dest}")
        # The plugin source lives beside this installer; warn if it (and so the
        # install source) resolves inside a worktree.
        ok_worktree = not _hook_bin_in_worktree(str(PLUGIN_SRC))
        if not ok_worktree:
            print(f"doctor [opencode]: WARNING: plugin source resolves inside a "
                  f"worktree ({PLUGIN_SRC}); re-install from the main checkout.")
        return 0 if (ok_ep and installed and ok_worktree) else 1


IBEGIN = "<!-- >>> data-olympus enforce (managed) >>> -->"
IEND = "<!-- <<< data-olympus enforce <<< -->"
INSTRUCTIONS_BODY = f"""{IBEGIN}
<!-- data-olympus-enforce v{SHIM_VERSION} -->
## Governed decisions (data-olympus)

Before making a code or architectural decision (choosing a library, pattern,
migration, API design, schema, or auth approach), you MUST call the `kb_consult`
MCP tool for the current workspace and follow the governing rules it returns.
Treat those rules as authoritative, not advisory. This agent cannot be hard-gated
locally, so compliance is observed via the data-olympus audit log.
{IEND}"""


class InstructionsProvider:
    """A soft provider that writes a MARKER-delimited managed block into a
    Markdown instructions file. Unlike the hard hook providers, it cannot block a
    tool call; it instructs the agent to consult the KB and relies on the
    data-olympus audit log to observe compliance. The block is delimited by the
    IBEGIN/IEND HTML comments so install is idempotent (one block, replaced on
    re-install) and uninstall is surgical (operator content preserved).
    """

    tier = "soft"

    def __init__(self, name: str, default_target: Path) -> None:
        self.name = name
        self._default_target = default_target

    def default_target(self) -> Path:
        return self._default_target

    @staticmethod
    def _strip_block(text: str) -> str:
        if IBEGIN in text and IEND in text:
            pre = text.split(IBEGIN, 1)[0].rstrip("\n")
            post = text.split(IEND, 1)[1].lstrip("\n")
            joined = "\n".join(p for p in (pre, post) if p)
            return (joined + "\n") if joined else ""
        return text

    def install(self, target: Path) -> int:
        existing = target.read_text() if target.exists() else ""
        if target.exists():
            _backup(target)
        base = self._strip_block(existing).rstrip("\n")
        new = (base + "\n\n" if base else "") + INSTRUCTIONS_BODY + "\n"
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(new)
        print(
            f"installed data-olympus enforcement (v{SHIM_VERSION}) into {target} "
            f"[{self.name}, tier=soft]"
        )
        return 0

    def uninstall(self, target: Path) -> int:
        if not target.exists():
            print("nothing to uninstall")
            return 0
        _backup(target)
        target.write_text(self._strip_block(target.read_text()))
        print(f"uninstalled data-olympus enforcement from {target} [{self.name}]")
        return 0

    def status(self, target: Path) -> int:
        if target.exists() and IBEGIN in target.read_text():
            print(f"{self.name}: installed, tier=soft, versions=['{SHIM_VERSION}']")
        else:
            print(f"{self.name}: not installed")
        return 0

    def doctor(self, _target: Path) -> int:
        ok, msg = _doctor_endpoint()
        print(f"doctor [{self.name}]: {msg}")
        return 0 if ok else 1


class UnsupportedProvider:
    """A documented-unsupported provider stub. The agent has no local hook,
    instructions, or MCP surface we can wire, so install/uninstall/doctor report
    the reason to stderr and exit non-zero (69 == EX_UNAVAILABLE); status reports
    unsupported but exits 0 (querying state is not itself a failure).
    """

    tier = "unsupported"

    def __init__(self, name: str, reason: str) -> None:
        self.name = name
        self._reason = reason

    def default_target(self) -> Path:
        return Path("/dev/null")

    def _report(self) -> int:
        print(f"{self.name}: unsupported -- {self._reason}", file=sys.stderr)
        return 69  # EX_UNAVAILABLE

    def install(self, _target: Path) -> int:
        return self._report()

    def uninstall(self, _target: Path) -> int:
        return self._report()

    def status(self, _target: Path) -> int:
        print(f"{self.name}: unsupported -- {self._reason}")
        return 0

    def doctor(self, _target: Path) -> int:
        return self._report()


HOOK_BEGIN = "# >>> data-olympus enforce (managed) >>>"
HOOK_END = "# <<< data-olympus enforce <<<"


def _git_managed_block(block: bool) -> str:
    if block:
        body = (
            'data-olympus report --staged --fail-on-unverified || {\n'
            '  echo "[KB] commit blocked: governed change without a consultation '
            '(set KB_ENFORCE_FAIL_MODE or run kb_consult)." >&2; exit 1;\n'
            '}'
        )
    else:
        body = (
            'data-olympus report --range HEAD~1..HEAD --emit-events || true'
        )
    return f"{HOOK_BEGIN}\n# data-olympus-enforce v{SHIM_VERSION}\n{body}\n{HOOK_END}\n"


class GitHookProvider:
    name = "git"
    tier = "detect"

    def __init__(self, *, block: bool = False) -> None:
        self._block = block

    def default_target(self) -> Path:
        hook = "pre-commit" if self._block else "post-commit"
        return Path(".git") / "hooks" / hook

    @staticmethod
    def _strip(text: str) -> str:
        if HOOK_BEGIN in text and HOOK_END in text:
            pre = text.split(HOOK_BEGIN, 1)[0].rstrip("\n")
            post = text.split(HOOK_END, 1)[1].lstrip("\n")
            joined = "\n".join(p for p in (pre, post) if p)
            return (joined + "\n") if joined else ""
        return text

    def install(self, target: Path) -> int:
        existing = target.read_text() if target.exists() else ""
        if target.exists():
            _backup(target)
        base = self._strip(existing).rstrip("\n")
        if not base:
            base = "#!/bin/sh"
        new = base + "\n\n" + _git_managed_block(self._block) + "\n"
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(new)
        target.chmod(0o755)
        kind = "pre-commit (block)" if self._block else "post-commit (warn)"
        print(f"installed data-olympus enforcement (v{SHIM_VERSION}) into {target} [git, {kind}]")
        return 0

    def uninstall(self, target: Path) -> int:
        if not target.exists():
            print("nothing to uninstall")
            return 0
        _backup(target)
        target.write_text(self._strip(target.read_text()))
        print(f"uninstalled data-olympus enforcement from {target} [git]")
        return 0

    def status(self, target: Path) -> int:
        if target.exists() and HOOK_BEGIN in target.read_text():
            print(f"git: installed, tier=detect, versions=['{SHIM_VERSION}']")
        else:
            print("git: not installed")
        return 0

    def doctor(self, target: Path) -> int:
        ok = target.exists() and HOOK_BEGIN in target.read_text() and os.access(target, os.X_OK)
        print(f"doctor [git]: hook {'present and executable' if ok else 'missing'} at {target}")
        return 0 if ok else 1


def registry() -> dict:
    return {
        "claude-code": _claude_provider(),
        "codex": _codex_provider(),
        "gemini": _gemini_provider(),
        "opencode": OpenCodeProvider(),
        # copilot-cli: the GitHub Copilot CLI loads "custom instructions from
        # AGENTS.md and related files" (verified via `copilot --help`:
        # `--no-custom-instructions` disables exactly that). No instructions file
        # exists in ~/.copilot/ on this laptop, so we default to the documented
        # global custom-instructions path ~/.copilot/copilot-instructions.md (one
        # of the "related files"). Operators can override with --settings.
        "copilot-cli": InstructionsProvider(
            "copilot-cli", Path(_home()) / ".copilot" / "copilot-instructions.md"),
        "copilot-ide": InstructionsProvider(
            "copilot-ide", Path(".github/copilot-instructions.md")),
        "antigravity": UnsupportedProvider(
            "antigravity",
            "no documented local hook/instructions/MCP surface as of 2026-06; "
            "revisit when Google publishes an extensibility API"),
    }


def main(argv: list[str]) -> int:
    p = argparse.ArgumentParser(prog="kb enforce")
    p.add_argument("command", choices=["install", "uninstall", "status", "doctor"])
    p.add_argument("--agent", default=None)
    p.add_argument("--all", action="store_true")
    p.add_argument("--settings", default=None)
    p.add_argument("--block", action="store_true",
                   help="git provider: install a blocking pre-commit hook "
                        "instead of post-commit warn")
    p.add_argument("--mode", choices=["off", "soft", "hard"], default="hard",
                   help="off uninstalls; soft installs inject-only (no blocking "
                        "gate); hard (default) installs the full gate")
    args = p.parse_args(argv)
    reg = registry()

    if args.all or (args.command == "status" and not args.agent and not args.settings):
        rc = 0
        for name, provider in reg.items():
            if args.all and getattr(provider, "tier", "") == "unsupported":
                print(f"{name}: skipped (unsupported)")
                continue
            cmd = args.command
            if args.command == "install" and args.mode == "off":
                cmd = "uninstall"
            elif args.command == "install" and args.mode == "soft":
                if hasattr(provider, "enforce_mode"):
                    provider.enforce_mode = "soft"
                else:
                    print(f"note: {name} has a fixed tier; --mode soft has no "
                          f"effect (use --mode off to uninstall)")
            target = provider.default_target()
            rc |= {
                "install": provider.install, "uninstall": provider.uninstall,
                "status": provider.status, "doctor": provider.doctor,
            }[cmd](target)
        return rc

    agent = args.agent or "claude-code"
    provider = GitHookProvider(block=args.block) if agent == "git" else reg.get(agent)
    if provider is None:
        print(f"kb enforce: unknown agent '{agent}' (known: {', '.join(sorted(reg))}, git)",
              file=sys.stderr)
        return 64
    command = args.command
    if args.command == "install" and args.mode == "off":
        command = "uninstall"
    elif args.command == "install" and args.mode == "soft":
        if hasattr(provider, "enforce_mode"):
            provider.enforce_mode = "soft"
        else:
            print(f"note: {agent} has a fixed tier; --mode soft has no effect "
                  f"(use --mode off to uninstall)")
    target = Path(args.settings) if args.settings else provider.default_target()
    return {
        "install": provider.install, "uninstall": provider.uninstall,
        "status": provider.status, "doctor": provider.doctor,
    }[command](target)


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
