#!/usr/bin/env python3
"""Headless CUBE notes agent — NO control panel. Reads a task from the CUBE
notes (a project), uses the cube note tools (note_search/note_list/
project_context + read_file), and writes a report. Designed to be driven by
cron (e.g. build an error-log report every hour); the interval is the cron
schedule, so it is fully configurable without touching the agent.

Usage: python3 agent.py --project NAME [--max-turns N] [--report PATH]
"""
import argparse, json, os, subprocess, sys, time
import datetime as _dt

def _now_ts():
    return _dt.datetime.now().strftime("%Y-%m-%d %H:%M")

MODEL = "/home/luulu/models/gemma-3-4b-it-Q4_K_M.gguf"
LLAMA = "/home/luulu/llama.cpp/build/bin/llama-server"
CUBE = "/home/CUBELinux/CUBELinux-2/target/release/cube"
PORT = 8080
BASE = f"http://127.0.0.1:{PORT}"
LOG = "/root/workspace/cube-agent/llama.log" if os.path.isdir("/root/workspace/cube-agent") else "/tmp/local-agent-llama.log"
TRACE_PROJECT = "agent-trace"   # every action the agent takes is logged here (replay/tuning)

def env():
    e = dict(os.environ)
    e["CUBE_NOTES_DIR"] = "/root/.cubelinux-notes"
    return e

def start_llama():
    p = subprocess.Popen(
        [LLAMA, "--model", MODEL, "--host", "127.0.0.1", "--port", str(PORT),
         "--alias", "gemma-3-4b", "--ctx-size", "8192", "-np", "1",
         "--temp", "0.2", "--repeat-penalty", "1.1"],
        stdout=open(LOG, "w"), stderr=subprocess.STDOUT, env=env())
    import requests
    for _ in range(180):
        try:
            if requests.get(BASE + "/v1/models", timeout=2).status_code == 200:
                break
        except Exception:
            pass
        time.sleep(2)
    return p

def chat(messages, max_tokens=600):
    import requests
    r = requests.post(BASE + "/v1/chat/completions", json={
        "model": "gemma-3-4b", "messages": messages,
        "max_tokens": max_tokens, "temperature": 0.08}, timeout=600)
    return r.json()["choices"][0]["message"]["content"]

def run_cube(args):
    try:
        r = subprocess.run([CUBE] + args, capture_output=True, text=True,
                           env=env(), timeout=30)
        return (r.stdout or r.stderr).strip()
    except Exception as e:
        return f"error: {e}"

def log_action(step, text):
    """Write an agent action to the CUBE notes trace (category=action) so the
    run can be replayed and fine-tuned."""
    body = f"step {step}: {text}"[:1600]
    subprocess.run([CUBE, "note", "add", "--project", TRACE_PROJECT,
                    "--cat", "action", body], capture_output=True, text=True, env=env())

def open_file(path):
    try:
        return open(path).read()[:2000]
    except Exception as e:
        return f"error: {e}"

# ---- incremental coverage (deterministic; offloaded from the small model) ----
import re as _re
COORD_RE = _re.compile(r'\b(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\b')

def _coords(text):
    return set(COORD_RE.findall(text))

# Bookkeeping categories that are NOT error-log findings (agent trace, coverage
# markers, resume checkpoints, test fixtures, design/impl notes, misc).
NONLOG_CATS = {"action", "checkpoint", "resume", "task", "debug", "impl", "design", "misc"}
# A finding-like title: a SHORT-UPPERCASE-HYPHEN/UNDERSCORE token then ':' or ' — '
FLAG_RE = _re.compile(r'\s*([A-Z][A-Z0-9_-]{2,})\s*[:—-]')
ERR_KW = ("error", "failed", "failure", "denied", "rejected", "cannot", "cannot be",
          "must ", "timeout", "times out", "not registered", "broken", "corrupt",
          "integrity", "permission denied", "selinux", "restorecon", "orphan",
          "eadrinuse", "shallow", "user_tmp_t", "mcp-client", "mcp plugin")

def _scan(cat, project=None):
    """Run `cube note list`, optionally filtered by category/project. With no cat
    and no project this scans the WHOLE store (every session), so the agent
    reviews ALL CUBE error logs, not one project or one category."""
    args = ["note", "list"]
    if cat:
        args += ["--cat", cat]
    if project:
        args += ["--project", project]
    return run_cube(args)

def _parse_note_line(line):
    """-> (coord, cat_or_None, subject) for a `note list` line, or None if the
    line is not a note (e.g. the 'N note(s):' header)."""
    m = _re.match(r'\s*(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\s+\[[^\]]+\]', line)
    if not m:
        return None
    coord = m.group(1)
    rest = line[m.end():].strip()
    cat = None
    cm = _re.match(r'\[([a-z]+)\]\s*', rest)
    if cm:
        cat = cm.group(1)
        rest = rest[cm.end():].strip()
    return coord, cat, rest

def _looks_like_finding(subj):
    """True if an un-categorized note looks like an error-log finding."""
    if FLAG_RE.match(subj):
        return True
    low = subj.lower()
    return any(k in low for k in ERR_KW)

def _log_findings(project=None):
    """[(coord, cat, subject)] for every finding-like note in the store, across ALL
    categories and projects. This includes `finding` notes and any other note with a
    FLAG-like / error title (the un-categorized logs), while excluding bookkeeping
    categories (action / checkpoint / resume / task / debug / impl / design / misc)."""
    out = []
    for line in _scan(None, project).splitlines():
        parsed = _parse_note_line(line)
        if parsed is None:
            continue
        coord, cat, subj = parsed
        if cat in NONLOG_CATS:
            continue
        if cat == "finding" or _looks_like_finding(subj):
            out.append((coord, cat, subj))
    return out

def _checkpoint_coords(project=None):
    """Own coord (leftmost) of every checkpoint note (optionally in a project)."""
    coords = []
    for line in _scan("checkpoint", project).splitlines():
        m = _re.match(r'\s*(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})', line)
        if m:
            coords.append(m.group(1))
    return coords

def _covered(project=None):
    """Covered = each checkpoint note's own coord + the coords in its FULL body.
    (note list only shows the truncated subject, so use note show to read the body.)"""
    out = set()
    for c in _checkpoint_coords(project):
        out.add(c)
        out |= _coords(run_cube(["note", "show", c]))
    return out

def uncovered_findings(project=None):
    covered = _covered(project)
    un = "".join(f"{coord}  {subj}\n" for coord, cat, subj in _log_findings(project) if coord not in covered)
    return un.strip() or "(no new findings)"

# Canonical issue key: collapse duplicate representations of the same root cause
# (e.g. FK1-SELINUX_RESTORECON vs SELINUX-RESTORECON) so the report dedupes by ISSUE,
# not just by coordinate. Prefixes (FK, FK#, FA, FB, FLAG, FIX) are stripped and
# separators normalized; a small alias map unifies differently-worded duplicates.
_ISSUE_PREFIX_RE = _re.compile(r'^(FK\d*|FA|FB|FLAG|FIX)[-_]?', _re.I)
_ISSUE_ALIAS = {
    "wal_checkpoint": "checkpoint_before_exit",
    "concurrent_store_checkpoint": "checkpoint_before_exit",
}

def _canonical_issue(flag):
    f = _ISSUE_PREFIX_RE.sub('', flag)
    f = _re.sub(r'[-_/\s]+', '_', f).strip('_').lower()
    return _ISSUE_ALIAS.get(f, f)

def _uncovered_meta(project=None):
    """[(coord, flag, desc)] for the currently-UNCOVERED findings, DEDUPED BY ISSUE.
    desc is the FULL note body; the most-detailed note per canonical issue is kept."""
    covered = _covered(project)
    best = {}   # issue -> (coord, flag, desc, len)
    for coord, cat, subj in _log_findings(project):
        if coord in covered:
            continue
        show = run_cube(["note", "show", coord])
        body_text = show.split("\n", 1)[1].strip() if "\n" in show else ""
        if not body_text:
            body_text = subj
        flag, sep, rest = body_text.partition(":")
        flag = flag.strip()
        desc = rest.strip() if sep else body_text
        issue = _canonical_issue(flag)
        cur = best.get(issue)
        if cur is None or len(desc) > cur[3]:
            best[issue] = (coord, flag, desc, len(desc))
    return [(c, f, d) for c, f, d, _ in best.values()]

# Snapshot of the uncovered findings captured when read_findings ran, so the
# report can be composed faithfully even after mark_covered has covered them.
CURRENT_META = []

def read_findings(project=None):
    """Deterministic, grounded input for the model: the currently-UNCOVERED CUBE
    findings, numbered with their FLAG. The model replies with ONLY the FLAGs in
    priority order; Python composes the report from the grounded text, so it is
    impossible for any finding to be invented."""
    global CURRENT_META
    meta = _uncovered_meta(project)
    CURRENT_META = meta
    if not meta:
        return "(no new findings)"
    lines = []
    for i, (coord, flag, desc) in enumerate(meta, 1):
        descf = (desc[:120] + "…") if len(desc) > 120 else desc
        lines.append(f"{i}. {flag}: {descf}  (coord {coord})")
    return ("Uncovered CUBE findings (use ONLY these; do not invent, add, or guess):\n"
            + "\n".join(lines)
            + "\n\nReply with the finding FLAGs in priority order, most critical first, "
              "comma-separated, using exactly the FLAGs above and nothing else — for example:\n"
            + ", ".join(f for _, f, _ in meta)
            + "\nThen call mark_covered, then reply done (no report body needed).")

_SEV_CRIT = ["data integrity", "security", "tls", "cert", "denied", "cannot",
             "corrupt", "lost", "inaccessible", "integrity", "breach"]
_SEV_HIGH = ["timeout", "persist", "reload", "shallow", "must", "fail", "broken"]

def _severity(flag, desc):
    """Deterministic severity label from the finding's FLAG + description. Used to
    give the report a meaningful priority order when the model provides no ranking."""
    t = (flag + " " + desc).lower()
    if any(k in t for k in _SEV_CRIT):
        return "HIGH"
    if any(k in t for k in _SEV_HIGH):
        return "MEDIUM"
    return "LOW"

def compose_report(project=None, model_text=""):
    """Build the grounded, prioritized report. The model supplies ONLY the priority
    ORDER of the FLAGs (honored when present); every finding's text is taken verbatim
    from CUBE, so any hallucinated finding is impossible. Unknown/duplicate flags are
    dropped; flags the model did not rank are ordered by a deterministic severity
    heuristic (HIGH → MEDIUM → LOW), then CUBE order. Uses the snapshot from read_findings."""
    meta = CURRENT_META or _uncovered_meta(project)
    if not meta:
        return "ERROR LOG REPORT (ALL)\n\nNo new findings since the last checkpoint."
    text_low = model_text.lower()
    listed = {}
    for coord, flag, desc in meta:
        listed.setdefault(flag.lower(), (coord, flag, desc))
    # Model's ranking first (by first-flag-appearance order in its reply).
    positions = []
    for key in listed:
        idx = text_low.find(key)
        if idx >= 0:
            positions.append((idx, key))
    positions.sort()
    order = [key for _, key in positions]
    remaining = [key for key in listed if key not in order]
    sev_rank = {"HIGH": 0, "MEDIUM": 1, "LOW": 2}
    remaining.sort(key=lambda k: sev_rank.get(_severity(listed[k][1], listed[k][2]), 3))
    order += remaining
    scope = project if project else "ALL"
    lines = [f"ERROR LOG REPORT (scope: {scope})",
             f"Generated: {_now_ts()}", ""]
    for rank, key in enumerate(order, 1):
        coord, flag, desc = listed[key]
        sev = _severity(flag, desc)
        lines.append(f"{rank}. [{sev}] {flag}: {desc if desc else '(no description)'}")
        lines.append(f"   coord: {coord}")
    return "\n".join(lines)

def mark_covered(project=None, text=""):
    covered = _covered(project)
    coords = [c for c, cat, subj in _log_findings(project) if c not in covered]
    if not coords:
        return "nothing to cover"
    body = "covered coords:\n" + "\n".join(coords)
    args = [CUBE, "note", "add", "--cat", "checkpoint"]
    if project:
        args += ["--project", project]
    args.append(body[:1400])
    subprocess.run(args, capture_output=True, text=True, env=env())
    return "marked covered"

def tool(kw, args, project):
    if kw == "read_findings":
        return read_findings(project)
    if kw == "mark_covered":
        return mark_covered(project, args)
    if kw == "note_search":
        return run_cube(["note", "search", args])
    if kw == "note_list":
        return run_cube(["note", "list"] + (args.split() if args else []))
    if kw == "project_context":
        return run_cube(["project", "context", args])
    if kw == "read_file":
        return open_file(args)
    return f"unknown tool: {kw}"

ACTIONS = ("read_findings", "mark_covered", "note_search",
           "note_list", "project_context", "read_file", "done")

def parse_action(text):
    """Parse an action word out of the model's reply. Tolerant of markdown code
    fences, ```tool_code``` wrappers, and prose preamble (the small gemma model
    tends to wrap calls like this). Scans lines and stops once it sees "done",
    so the trailing report is never misread as another tool call."""
    best = (None, "")
    for raw in text.strip().splitlines():
        s = raw.replace("`", "").replace("tool_code", " ").strip()
        low = s.lower()
        for kw in ACTIONS:
            idx = low.find(kw)
            if idx >= 0:
                best = (kw, s[idx + len(kw):].strip())
                break
        if best[0] == "done":
            break
    return best

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--project", default=None,
                    help="Restrict the report to one CUBE project (default: ALL finding notes across the store)")
    ap.add_argument("--max-turns", type=int, default=6)
    ap.add_argument("--report", default="/root/workspace/cube-agent/report.txt")
    a = ap.parse_args()

    p = start_llama()
    scope = a.project if a.project else "ALL"
    # Ground the report in CUBE up front: snapshot the uncovered findings and hand
    # them to the model so it only has to RANK them (it never writes finding text).
    findings_text = read_findings(a.project)
    if findings_text == "(no new findings)":
        initial_context = "There are no uncovered findings to report right now."
        task = "Reply done (no report body needed)."
    else:
        initial_context = findings_text
        task = ("Reply with the finding FLAGs in priority order (most critical first), "
                "comma-separated, using exactly the FLAGs below and nothing else — then reply done.")
    sys_prompt = (
        f"You are a headless error-log reviewer for the CUBE notes store (scope: {scope}).\n"
        "The findings are provided below. Order them by priority/severity, most critical first.\n"
        "Reply with the FLAGs in that order, comma-separated, using exactly those FLAGs and nothing else.\n"
        "Then reply: done. "
        "Your ranked FLAG list sets the priority; the report body is assembled for you, so do not invent findings."
    )
    messages = [
        {"role": "system", "content": sys_prompt},
        {"role": "user", "content": f"{task}\n\n{initial_context}"},
    ]
    trace = []
    report_text = ""
    last_assistant = ""
    for turn in range(a.max_turns):
        try:
            text = chat(messages)
        except Exception as e:
            trace.append(f"chat error: {e}")
            break
        kw, args = parse_action(text)
        if kw == "done":
            # The model supplies ONLY the priority order of the FLAGs; Python
            # composes the grounded report (impossible to invent a finding).
            di = text.lower().find("done")
            model_text = text[di + 4:].strip() if di >= 0 else last_assistant
            report_text = compose_report(a.project, model_text)
            log_action(turn, f"DONE\n\n~flags: {model_text[:400]}")
            break
        if kw:
            result = tool(kw, args, a.project)
            trace.append(f"turn{turn}: {kw}({args})")
            log_action(turn, f"{kw}({args})\n\n{result[:1400]}")
            messages.append({"role": "assistant", "content": text})
            messages.append({"role": "user",
                             "content": f"Tool result:\n{result}\n\nContinue — output exactly one "
                                         "action line (read_findings / mark_covered / note_search / "
                                         "note_list / project_context / read_file / done)."})
        else:
            log_action(turn, f"(no action) {text[:300]}")
            messages.append({"role": "assistant", "content": text})
            messages.append({"role": "user",
                             "content": "Output exactly ONE action line starting with a tool name."})
            trace.append(f"turn{turn}: (no action)")
        last_assistant = text
    if not report_text:
        # Run ended without an explicit "done" — still produce a grounded report.
        report_text = compose_report(a.project, last_assistant)
    # Deterministic coverage: mark all uncovered findings covered (including any
    # deduped-away duplicates) regardless of whether the model did it, so the next
    # run is genuinely incremental. No-op when there are no uncovered findings.
    mark_covered(a.project)
    with open(a.report, "w") as f:
        f.write(report_text)
    os.makedirs(os.path.dirname(a.report), exist_ok=True) if os.path.dirname(a.report) else None
    print("AGENT_DONE max_turns=", a.max_turns, "report=", a.report,
          "bytes=", os.path.getsize(a.report) if os.path.exists(a.report) else 0)
    try:
        p.terminate()
    except Exception:
        pass

if __name__ == "__main__":
    main()