#!/usr/bin/env python3
"""Structural guardrail: every write to an observation-determining field of the
daemon Hub's clientState must run under the machine observation guard write
side.

This is a block-aware (not a line/brace regex) checker: it parses protocol.go's
functions with proper brace nesting and, for every mutation of
sessionGeneration/agentSessions/agentTurnStateRecords/turnStateHighWater/
connectedAt, proves the enclosing function either

  (a) lexically contains a guard-write acquisition (acquireWrite /
      withObservationGuardWrite) whose scope still covers the mutation line, or
  (b) is a `*StateLocked` helper that is only ever called from guarded
      functions (so it can never execute outside the guard boundary).

A mutation placed BEFORE the guard acquisition, or after the matching release,
or inside a helper that is also called from an unguarded function, fails. The
old checker accepted `client.sessionGeneration = ...` at the start of an
allowlisted function; this one requires the guard acquisition to lexically
dominate the mutation.
"""

import pathlib
import re
import sys

ROOT = pathlib.Path(__file__).resolve().parent.parent.parent
PROTOCOL = ROOT / "internal/api/daemon/protocol.go"

OBSERVATION_FIELDS = [
    "sessionGeneration",
    "agentSessions[",
    "agentTurnStateRecords[",
    "turnStateHighWater",
    "connectedAt",
]

# Functions that are NOT StateLocked helpers and NOT themselves guarded, but
# legitimately own the guard because they acquire the write side inline and the
# checker sees the acquisition in the same function.
GUARD_ACQUIRE_TOKENS = ["acquireWrite()", "withObservationGuardWrite("]


class FuncInfo:
    def __init__(self, name, start, end, body_start):
        self.name = name
        self.start = start
        self.end = end
        self.body_start = body_start  # line after the opening brace
        self.acquire_lines = []  # lines where a guard write is acquired
        self.release_lines = []  # lines where the guard write is released


def parse_functions(lines):
    """Returns a list of FuncInfo using proper brace matching.

    For a `func ... X(...) {` declaration, the body starts on the line AFTER
    the opening brace and ends at the line whose brace depth returns to zero.
    Multi-line signatures and comments with braces are tolerated by net-depth
    counting across the whole file from the decl line.
    """
    funcs = []
    i = 0
    n = len(lines)
    while i < n:
        line = lines[i]
        m = re.match(r"^func (\([^)]*\) )?([A-Za-z_][A-Za-z0-9_]*)\(", line)
        if not m:
            i += 1
            continue
        name = m.group(2)
        # Find the opening brace of the body: scan forward tracking net depth
        # until the depth first becomes positive (the `{` that opens the body).
        depth = 0
        j = i
        while j < n:
            depth += lines[j].count("{") - lines[j].count("}")
            if depth > 0:
                break
            j += 1
        if j >= n:
            i += 1
            continue
        # body_start = the line AFTER the opening brace.
        body_start = j + 1
        # Scan forward to the matching close: depth returns to zero.
        depth = 0
        k = j
        while k < n:
            depth += lines[k].count("{") - lines[k].count("}")
            if depth == 0 and k > j:
                break
            k += 1
        funcs.append(FuncInfo(name, i, k, body_start))
        i = k + 1
    return funcs


def is_comment(line):
    stripped = line.strip()
    return stripped.startswith("//") or stripped.startswith("/*") or stripped.startswith("*")


def is_write_to_field(line, field):
    if is_comment(line):
        return False
    if "delete(" in line or "clear(" in line:
        return field in line
    if "=" not in line:
        return False
    lhs, _, _ = line.partition("=")
    normalized = lhs.replace(" ", "")
    # Actual assignment: the field access must be on the LHS. A line like
    # `if strings.TrimSpace(client.agentSessions[agentID]) == "" {` has the field
    # on the LHS of `==`, not `=`, so partition("=") splits at the FIRST `=` which
    # is the `==` comparison — normalized would include the field and misjudge.
    # Require the LHS to START with the client field access (an assignment) or be
    # a compound like `client.agentSessions[k] = `.
    if normalized.startswith("client." + field):
        return True
    # delete/clear handled above; also allow `client.agentSessions[k] = `.
    if ("client." + field) in normalized and "==" not in line:
        return True
    return False


GUARD_ACQUIRE_RE = re.compile(r"\.acquireWrite\(\)|withObservationGuardWrite\(")
GUARD_RELEASE_RE = re.compile(r"\.releaseWrite\(\)")


def guarded_scope(lines, func, line_no):
    """Returns True when the write at line_no is lexically inside a guard-write
    acquisition within this function: a acquireWrite()/withObservationGuardWrite(
    appears at a line before line_no and its matching releaseWrite()/end of the
    callback appears after line_no, within the function body."""
    # withObservationGuardWrite(func() { ... }) — the mutation inside the closure
    # is covered by the callback. acquireWrite() ... releaseWrite() — the
    # mutation between them is covered.
    acquire_line = None
    release_line = None
    deferred_release = False
    for i in range(func.body_start, line_no):
        line = lines[i]
        if GUARD_ACQUIRE_RE.search(line):
            acquire_line = i
        if GUARD_RELEASE_RE.search(line):
            release_line = i
            if line.strip().startswith("defer"):
                deferred_release = True
    if acquire_line is None:
        return False
    # The inline acquireWrite()/releaseWrite() form: a release AFTER the
    # mutation (within the function) proves the write is inside the guarded
    # scope, even if the release has not appeared yet before the mutation.
    if release_line is not None and release_line > line_no:
        return True
    # A `defer client.guard.releaseWrite()` before the mutation means the write
    # side is HELD for the remainder of the function — any later mutation is
    # guarded.
    if deferred_release:
        return True
    # If no release precedes the mutation, scan FORWARD for a matching release
    # inside the same function (inline form) — the mutation between acquire and
    # release is guarded.
    if release_line is None:
        for i in range(line_no, func.end + 1):
            if GUARD_RELEASE_RE.search(lines[i]):
                return True
    # Callback form: a `})` closing the withObservationGuardWrite must appear
    # after the mutation.
    for i in range(line_no + 1, func.end + 1):
        if "})" in lines[i] or "});" in lines[i]:
            return True
        # A nested function closing could terminate before the callback end.
        if re.match(r"^\s*}\s*\)\s*$", lines[i]):
            return True
    return False


def called_from_unguarded(funcs, name, call_graph, guarded_funcs):
    """StateLocked helpers may only be called from guarded functions."""
    from collections import deque

    # Build reverse call graph.
    callers = {}
    for caller, callees in call_graph.items():
        for callee in callees:
            callers.setdefault(callee, set()).add(caller)

    seen = set()
    queue = deque([name])
    while queue:
        current = queue.popleft()
        if current in seen:
            continue
        seen.add(current)
        for caller in callers.get(current, set()):
            if caller in guarded_funcs:
                continue
            queue.append(caller)
    # If any caller is reachable that is not guarded, the helper is unsafe.
    for caller in callers.get(name, set()):
        if caller not in guarded_funcs:
            return True
    return False


def main():
    src = PROTOCOL.read_text()
    lines = src.splitlines()
    funcs = parse_functions(lines)

    # Build a simple call graph from identifier calls.
    call_graph = {}
    for func in funcs:
        callees = set()
        for i in range(func.body_start, func.end + 1):
            for m in re.finditer(r"client\.([A-Za-z_][A-Za-z0-9_]*)\(|([A-Za-z_][A-Za-z0-9_]*)\(\)", lines[i]):
                callees.add(m.group(1) or m.group(2))
        call_graph[func.name] = callees

    # Determine which functions lexically acquire the guard write side.
    guarded_funcs = set()
    for func in funcs:
        for i in range(func.body_start, func.end + 1):
            if GUARD_ACQUIRE_RE.search(lines[i]):
                guarded_funcs.add(func.name)
                break

    violations = []
    for func in funcs:
        for i in range(func.body_start, func.end + 1):
            if is_comment(lines[i]):
                continue
            for field in OBSERVATION_FIELDS:
                if not is_write_to_field(lines[i], field):
                    continue
                # Accept if lexically inside a guard write scope.
                if guarded_scope(lines, func, i):
                    break
                # Accept guarded-write helpers: *StateLocked / *Locked helpers
                # that are only called from guarded functions can never execute
                # outside the guard boundary.
                if (func.name.endswith("StateLocked") or func.name.endswith("Locked")) and not called_from_unguarded(funcs, func.name, call_graph, guarded_funcs):
                    break
                violations.append((i + 1, func.name, field, lines[i].strip()))
                break

    if violations:
        print(f"FAIL: {len(violations)} unguarded observation-field write(s)")
        for line_no, fn, field, text in violations[:40]:
            print(f"  {PROTOCOL}:{line_no} [{fn}] {field}: {text}")
        print("Every write to sessionGeneration/agentSessions/agentTurnStateRecords/turnStateHighWater/connectedAt must run under the machine observation guard write side.")
        return 1
    print("Guarded writer guardrail passed.")
    return 0


if __name__ == "__main__":
    sys.exit(main())