from __future__ import annotations

import errno
import fcntl
import json
import os
import re
import stat
import tempfile
from collections.abc import Callable, Iterable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

ACCOUNT_LOCKS_FILENAME = "account_locks.json"
ACCOUNT_LOCK_SCHEMA = "account-lock/v1"
ACCOUNT_LOCK_EXIT_CODE = 77
ACCOUNT_LOCK_MAX_BYTES = 64 * 1024
ACCOUNT_LOCK_TOOLS = frozenset({"codex", "claude", "grok"})

# Mirror modules/workstation/claude/lib/session-class.sh: agent/dispatch ancestry
# is decisive, human.slice is affirmative, and an ambiguous launch needs a TTY.
_AGENT_CGROUP_RE = re.compile(r"/agent\.slice(?:/|$)")
_DISPATCH_CGROUP_RE = re.compile(r"/[^/]*[Dd]ispatch[^/]*\.scope(?:/|$)")
_HUMAN_CGROUP_RE = re.compile(r"/human\.slice(?:/|$)")
_CONTAINER_CGROUP_RE = re.compile(
    r"(?:libpod|docker|containerd|/machine\.slice/|kubepods)", re.IGNORECASE
)
_CGROUP_ANCESTRY_MAX = 64


class AccountLockError(RuntimeError):
    """Base refusal for account-lock policy and state failures."""

    exit_code = ACCOUNT_LOCK_EXIT_CODE
    reason = "account-lock-refused"

    def refusal_payload(self) -> dict[str, object]:
        return {
            "ok": False,
            "detail": "account-locked",
            "reason": self.reason,
        }


class AccountLockStateError(AccountLockError):
    reason = "account-lock-state-invalid"

    def __init__(self, path: Path, detail: str) -> None:
        self.path = Path(path)
        super().__init__(f"account lock state is invalid at {self.path}: {detail}")

    def refusal_payload(self) -> dict[str, object]:
        payload = super().refusal_payload()
        payload["detail"] = "account-lock-state-invalid"
        return payload


class AccountLockedError(AccountLockError):
    reason = "account-locked"

    def __init__(self, tool: str, slugs: Iterable[str], *, command: str | None = None) -> None:
        normalized = tuple(dict.fromkeys(slug for slug in slugs if slug))
        self.tool = tool
        self.slugs = normalized
        account_text = ", ".join(f"{tool}:{slug}" for slug in normalized) or tool
        prefix = f"{command}: " if command else ""
        super().__init__(
            f"{prefix}account lock blocks {account_text}; unlock it in Systray AI or use "
            "--human-override-lock with an explicit --account from a human terminal"
        )

    def refusal_payload(self) -> dict[str, object]:
        payload = super().refusal_payload()
        payload["accounts"] = [f"{self.tool}:{slug}" for slug in self.slugs]
        return payload


class HumanOverrideError(AccountLockError):
    reason = "human-override-denied"

    def refusal_payload(self) -> dict[str, object]:
        payload = super().refusal_payload()
        payload["detail"] = "human-override-denied"
        return payload


@dataclass(frozen=True, slots=True)
class HumanSessionVerdict:
    human: bool
    reason: str


@dataclass(frozen=True, slots=True)
class AccountUseAuthorization:
    tool: str
    slug: str
    locked_override: bool

    @property
    def key(self) -> str:
        return account_key(self.tool, self.slug)


def account_key(tool: str, slug: str) -> str:
    normalized_tool = tool.strip()
    normalized_slug = slug.strip()
    if normalized_tool not in ACCOUNT_LOCK_TOOLS:
        raise ValueError(f"unsupported account-lock tool: {tool!r}")
    if (
        not normalized_slug
        or ":" in normalized_slug
        or "/" in normalized_slug
        or "\\" in normalized_slug
        or any(ord(char) < 32 or ord(char) == 127 for char in normalized_slug)
    ):
        raise ValueError(f"invalid account slug for account lock: {slug!r}")
    return f"{normalized_tool}:{normalized_slug}"


def split_account_key(key: str) -> tuple[str, str]:
    if not isinstance(key, str) or ":" not in key:
        raise ValueError("account lock entries must use tool:slug keys")
    tool, slug = key.split(":", 1)
    account_key(tool, slug)
    return tool, slug


class AccountLockStore:
    """Persistent shared-tier admission state for new account-backed sessions."""

    def __init__(self, base_dir: Path) -> None:
        self.base_dir = Path(base_dir)
        self.path = self.base_dir / ACCOUNT_LOCKS_FILENAME
        self.lock_path = self.base_dir / ".account_locks.lock"

    def locked_keys(self) -> frozenset[str]:
        with self._state_lock():
            return frozenset(self._read_unlocked())

    def locked_slugs(self, tool: str) -> frozenset[str]:
        if tool not in ACCOUNT_LOCK_TOOLS:
            raise ValueError(f"unsupported account-lock tool: {tool!r}")
        prefix = f"{tool}:"
        return frozenset(
            key[len(prefix) :] for key in self.locked_keys() if key.startswith(prefix)
        )

    def is_locked(self, tool: str, slug: str) -> bool:
        return account_key(tool, slug) in self.locked_keys()

    def set_locked(self, tool: str, slug: str, locked: bool) -> bool:
        key = account_key(tool, slug)
        with self._state_lock():
            keys = self._read_unlocked()
            changed = (key not in keys) if locked else (key in keys)
            if not changed:
                return False
            if locked:
                keys.add(key)
            else:
                keys.remove(key)
            self._write_unlocked(keys)
            return True

    def rekey(self, tool: str, old_slug: str, new_slug: str) -> bool:
        old_key = account_key(tool, old_slug)
        new_key = account_key(tool, new_slug)
        with self._state_lock():
            keys = self._read_unlocked()
            if old_key not in keys:
                return False
            keys.remove(old_key)
            keys.add(new_key)
            self._write_unlocked(keys)
            return True

    def drop(self, tool: str, slug: str) -> bool:
        return self.set_locked(tool, slug, False)

    def restore_key_states(self, states: Mapping[str, bool]) -> bool:
        """Restore selected keys without clobbering unrelated concurrent updates."""

        normalized = {
            account_key(*split_account_key(key)): bool(locked)
            for key, locked in states.items()
        }
        with self._state_lock():
            keys = self._read_unlocked()
            before = set(keys)
            for key, locked in normalized.items():
                if locked:
                    keys.add(key)
                else:
                    keys.discard(key)
            if keys == before:
                return False
            self._write_unlocked(keys)
            return True

    @contextmanager
    def _state_lock(self) -> Iterator[None]:
        try:
            self.base_dir.mkdir(parents=True, exist_ok=True)
        except OSError as exc:
            raise AccountLockStateError(
                self.base_dir, f"cannot create account-lock directory: {exc}"
            ) from exc
        flags = os.O_RDWR | os.O_CREAT
        if hasattr(os, "O_CLOEXEC"):
            flags |= os.O_CLOEXEC
        if hasattr(os, "O_NOFOLLOW"):
            flags |= os.O_NOFOLLOW
        try:
            fd = os.open(self.lock_path, flags, 0o600)
        except OSError as exc:
            detail = "lock path must not be a symlink" if exc.errno == errno.ELOOP else str(exc)
            raise AccountLockStateError(self.lock_path, detail) from exc
        try:
            lock_stat = os.fstat(fd)
            if not stat.S_ISREG(lock_stat.st_mode):
                raise AccountLockStateError(self.lock_path, "lock path is not a regular file")
            os.fchmod(fd, 0o600)
            fcntl.flock(fd, fcntl.LOCK_EX)
            try:
                yield
            finally:
                fcntl.flock(fd, fcntl.LOCK_UN)
        finally:
            os.close(fd)

    def _read_unlocked(self) -> set[str]:
        flags = os.O_RDONLY
        if hasattr(os, "O_CLOEXEC"):
            flags |= os.O_CLOEXEC
        if hasattr(os, "O_NOFOLLOW"):
            flags |= os.O_NOFOLLOW
        try:
            fd = os.open(self.path, flags)
        except FileNotFoundError:
            return set()
        except OSError as exc:
            detail = "state path must not be a symlink" if exc.errno == errno.ELOOP else str(exc)
            raise AccountLockStateError(self.path, detail) from exc

        try:
            state_stat = os.fstat(fd)
            if not stat.S_ISREG(state_stat.st_mode):
                raise AccountLockStateError(self.path, "state path is not a regular file")
            if state_stat.st_size > ACCOUNT_LOCK_MAX_BYTES:
                raise AccountLockStateError(self.path, "state file is too large")
            if state_stat.st_mode & (stat.S_IWGRP | stat.S_IWOTH):
                raise AccountLockStateError(self.path, "state file is group/world writable")
            chunks: list[bytes] = []
            remaining = ACCOUNT_LOCK_MAX_BYTES + 1
            while remaining > 0:
                chunk = os.read(fd, min(remaining, 8192))
                if not chunk:
                    break
                chunks.append(chunk)
                remaining -= len(chunk)
            raw = b"".join(chunks)
        finally:
            os.close(fd)

        if len(raw) > ACCOUNT_LOCK_MAX_BYTES:
            raise AccountLockStateError(self.path, "state file is too large")
        try:
            payload = json.loads(raw.decode("utf-8", errors="strict"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise AccountLockStateError(self.path, "state file is not valid UTF-8 JSON") from exc
        if not isinstance(payload, dict):
            raise AccountLockStateError(self.path, "state root must be an object")
        if payload.get("version") != ACCOUNT_LOCK_SCHEMA:
            raise AccountLockStateError(
                self.path, f"version must be {ACCOUNT_LOCK_SCHEMA!r}"
            )
        locked = payload.get("locked")
        if not isinstance(locked, list):
            raise AccountLockStateError(self.path, "locked must be a list")
        normalized: set[str] = set()
        try:
            for key in locked:
                if not isinstance(key, str):
                    raise ValueError("account lock entries must be strings")
                normalized.add(account_key(*split_account_key(key)))
        except ValueError as exc:
            raise AccountLockStateError(self.path, str(exc)) from exc
        if len(normalized) != len(locked):
            raise AccountLockStateError(self.path, "locked contains duplicate entries")
        return normalized

    def _write_unlocked(self, keys: Iterable[str]) -> None:
        normalized = sorted({account_key(*split_account_key(key)) for key in keys})
        payload = {
            "version": ACCOUNT_LOCK_SCHEMA,
            "locked": normalized,
        }
        self.base_dir.mkdir(parents=True, exist_ok=True)
        temp_path: Path | None = None
        try:
            with tempfile.NamedTemporaryFile(
                "w",
                encoding="utf-8",
                dir=self.base_dir,
                prefix=".account_locks.",
                suffix=".tmp",
                delete=False,
            ) as handle:
                json.dump(payload, handle, indent=2, sort_keys=True)
                handle.write("\n")
                handle.flush()
                os.fsync(handle.fileno())
                os.fchmod(handle.fileno(), 0o600)
                temp_path = Path(handle.name)
            os.replace(temp_path, self.path)
            directory_fd = os.open(self.base_dir, os.O_RDONLY)
            try:
                os.fsync(directory_fd)
            finally:
                os.close(directory_fd)
        except OSError as exc:
            raise AccountLockStateError(self.path, f"cannot write state: {exc}") from exc
        finally:
            if temp_path is not None:
                try:
                    if temp_path.exists() or temp_path.is_symlink():
                        temp_path.unlink()
                except OSError:
                    pass


def classify_human_session(
    cgroup_paths: Iterable[str],
    *,
    controlling_tty: bool,
    container_signal: str | None,
) -> HumanSessionVerdict:
    paths = tuple(path for path in cgroup_paths if path)
    for path in paths:
        if _AGENT_CGROUP_RE.search(path) or _DISPATCH_CGROUP_RE.search(path):
            return HumanSessionVerdict(False, f"agent-cgroup:{path}")
    if container_signal is not None:
        return HumanSessionVerdict(False, f"in-container:{container_signal}")
    for path in paths:
        if _HUMAN_CGROUP_RE.search(path):
            return HumanSessionVerdict(True, f"human-cgroup:{path}")
    if controlling_tty:
        return HumanSessionVerdict(True, "controlling-tty")
    return HumanSessionVerdict(False, "no-tty")


def probe_human_session(pid: int | None = None) -> HumanSessionVerdict:
    target_pid = os.getpid() if pid is None else pid
    cgroups = tuple(_cgroup_ancestry(target_pid))
    return classify_human_session(
        cgroups,
        controlling_tty=_has_controlling_tty(target_pid),
        container_signal=_container_signal(cgroups),
    )


def _cgroup_ancestry(pid: int) -> Iterator[str]:
    current = pid
    for _depth in range(_CGROUP_ANCESTRY_MAX):
        if current <= 1:
            return
        cgroup = _cgroup_of(current)
        if cgroup is not None:
            yield cgroup
        parent = _ppid_of(current)
        if parent is None or parent == current:
            return
        current = parent


def _cgroup_of(pid: int) -> str | None:
    try:
        line = (Path("/proc") / str(pid) / "cgroup").read_text(encoding="utf-8").splitlines()[0]
    except (OSError, IndexError, UnicodeError):
        return None
    if line.startswith("0::"):
        return line[3:]
    parts = line.split(":", 2)
    return parts[2] if len(parts) == 3 else None


def _proc_stat_tail(pid: int) -> list[str] | None:
    try:
        raw = (Path("/proc") / str(pid) / "stat").read_text(encoding="utf-8")
    except (OSError, UnicodeError):
        return None
    marker = raw.rfind(") ")
    if marker < 0:
        return None
    return raw[marker + 2 :].split()


def _ppid_of(pid: int) -> int | None:
    fields = _proc_stat_tail(pid)
    if fields is None or len(fields) < 2:
        return None
    try:
        return int(fields[1])
    except ValueError:
        return None


def _has_controlling_tty(pid: int) -> bool:
    fields = _proc_stat_tail(pid)
    if fields is None or len(fields) < 5:
        return False
    try:
        return int(fields[4]) != 0
    except ValueError:
        return False


def _container_signal(cgroups: Iterable[str]) -> str | None:
    if Path("/run/.containerenv").exists():
        return "file:/run/.containerenv"
    if Path("/.dockerenv").exists():
        return "file:/.dockerenv"
    container = os.environ.get("container")
    if container:
        return f"env:container={container}"
    if any(_CONTAINER_CGROUP_RE.search(path) for path in cgroups):
        return "cgroup:/proc ancestry"
    return None


def confirmation_phrase(
    action: Literal["lock", "unlock", "use-once"], tool: str, slug: str
) -> str:
    key = account_key(tool, slug)
    if action == "lock":
        return f"LOCK {key}"
    if action == "unlock":
        return f"UNLOCK {key}"
    if action == "use-once":
        return f"USE {key} ONCE"
    raise ValueError(f"unsupported account lock action: {action}")


def confirm_human_action(
    action: Literal["lock", "unlock", "use-once"],
    tool: str,
    slug: str,
    *,
    session_probe: Callable[[], HumanSessionVerdict] = probe_human_session,
    confirmation_reader: Callable[[str], str] | None = None,
) -> None:
    verdict = session_probe()
    if not verdict.human:
        raise HumanOverrideError(
            f"account lock action requires a human terminal; denied ({verdict.reason})"
        )
    phrase = confirmation_phrase(action, tool, slug)
    prompt = f"Type {phrase!r} to confirm: "
    reader = confirmation_reader or _read_confirmation_from_tty
    try:
        entered = reader(prompt)
    except OSError as exc:
        raise HumanOverrideError(
            f"account lock action requires /dev/tty confirmation: {exc}"
        ) from exc
    if entered.strip() != phrase:
        raise HumanOverrideError("account lock confirmation did not match")


def _read_confirmation_from_tty(prompt: str) -> str:
    with Path("/dev/tty").open("r+", encoding="utf-8", buffering=1) as terminal:
        terminal.write(prompt)
        terminal.flush()
        return terminal.readline(512)


def authorize_account_use(
    base_dir: Path,
    tool: str,
    slug: str,
    *,
    allow_locked: bool,
    explicit_account: bool,
    command: str | None = None,
) -> AccountUseAuthorization:
    store = AccountLockStore(base_dir)
    if not store.is_locked(tool, slug):
        return AccountUseAuthorization(tool, slug, False)
    if not allow_locked:
        raise AccountLockedError(tool, (slug,), command=command)
    if not explicit_account:
        prefix = f"{command}: " if command else ""
        raise HumanOverrideError(
            f"{prefix}--human-override-lock requires an explicit --account"
        )
    confirm_human_action("use-once", tool, slug)
    return AccountUseAuthorization(tool, slug, True)


def revalidate_account_use(base_dir: Path, authorization: AccountUseAuthorization) -> None:
    locked = AccountLockStore(base_dir).is_locked(
        authorization.tool, authorization.slug
    )
    if authorization.locked_override:
        return
    if locked:
        raise AccountLockedError(authorization.tool, (authorization.slug,))