"""Boardroom Map role agent — a sandboxed, ONE-SHOT, single-completion agent.

One container = one role = one model call (plus a bounded reliability ladder /
repair round-trip). No tools, no loops, no shared writable workspace. The
orchestrator (orchestrator/graders.py + adjudicator.py) launches this hardened
(non-root, read-only rootfs, per-job network) with BM_ROLE set to:

    extractor    reads /docs (ro) + the BDEF red-flag taxonomy, emits STRUCTURED
                 JSON per /schema.json (extraction schema) -> /out/extraction.json
    grader       reads /docs (ro) + the full /BDEF.md rubric, scores categories
                 A-H per /schema.json (grades schema)      -> /out/<BM_GRADER_ID>.json
    adjudicator  reads /extraction.json (ro) + the panel's /grades/*.json (ro),
                 writes a MARKDOWN adjudication (no scores) -> /out/ADJUDICATION.md

Env (set by the orchestrator):
    BM_ROLE, BM_GRADER_ID, BM_GRADER_NAME, BM_MODEL,
    BM_LLM_BASE (http://boardroom-proxy:4000/v1), BM_LLM_KEY,
    BM_TEMPERATURE (extractor forced to 0.0), BM_MAX_MODEL_LEN

Mounts: /docs (ro), /BDEF.md (ro), /schema.json (ro; role-appropriate),
/persona/PERSONA.md (ro, optional), /out (rw); adjudicator additionally
/grades (ro) and /extraction.json (ro).

JSON reliability ladder (extractor + grader):
    1. response_format = {"type":"json_schema", ..., "strict": true}
    2. on HTTP 4xx: retry with top-level {"guided_json": <schema>} (vLLM ext.)
    3. on another 4xx: retry plain
Parse whole-reply JSON, else the first brace-balanced {...} block. If invalid,
ONE repair round-trip; if still bad, write the raw text to the output path plus
a sibling <output>.invalid marker containing the error. Full jsonschema
validation runs orchestrator-side; only lightweight structural checks here.

Air-gapped: the per-job Docker network is --internal, so the only reachable
endpoint is the model proxy. Pure Python stdlib (urllib) — no pip deps.
"""
from __future__ import annotations

import json
import os
import ssl
import time
import traceback
import urllib.error
import urllib.request

_SSL_CTX = ssl._create_unverified_context()  # LAN self-signed certs

ROLE = os.environ.get("BM_ROLE", "grader")  # extractor | grader | adjudicator
GID = os.environ.get("BM_GRADER_ID", ROLE)
NAME = os.environ.get("BM_GRADER_NAME", GID)
MODEL = os.environ.get("BM_MODEL", "grader-a")
LLM_BASE = os.environ.get("BM_LLM_BASE", "http://boardroom-proxy:4000/v1").rstrip("/")
LLM_KEY = os.environ.get("BM_LLM_KEY", "sk-local")
try:
    TEMPERATURE = float(os.environ.get("BM_TEMPERATURE", "") or "0.3")
except ValueError:
    TEMPERATURE = 0.3
try:
    MAX_MODEL_LEN = int(os.environ.get("BM_MAX_MODEL_LEN", "32768"))
except ValueError:
    MAX_MODEL_LEN = 32768

DOCS = "/docs"
BDEF_PATH = "/BDEF.md"
SCHEMA_PATH = "/schema.json"
PERSONA_PATH = "/persona/PERSONA.md"
GRADES_DIR = "/grades"
EXTRACTION_PATH = "/extraction.json"
OUT_DIR = "/out"

MAX_OUTPUT_TOKENS = 3072
TRANSPORT_RETRIES = 4


def log(msg: str) -> None:
    print(f"[{GID}] {msg}", flush=True)


# ---------------------------------------------------------------- io helpers
def read_text(path: str, limit: int = 2_000_000) -> str:
    try:
        with open(path, errors="replace") as f:
            return f.read()[:limit]
    except (FileNotFoundError, IsADirectoryError):
        return ""


def _doc_names() -> list[str]:
    if not os.path.isdir(DOCS):
        return []
    return sorted(fn for fn in os.listdir(DOCS)
                  if fn.endswith(".txt") and os.path.isfile(os.path.join(DOCS, fn)))


def preload_docs(budget: int) -> tuple[str, bool]:
    """Concatenate /docs/*.txt up to `budget` chars. Returns (text, truncated)."""
    chunks, used, truncated = [], 0, False
    for n in _doc_names():
        body = read_text(os.path.join(DOCS, n))
        header = f"\n\n========== DOCUMENT: {n} ==========\n"
        room = budget - used
        if room <= 0:
            truncated = True
            break
        seg = (header + body)[:room]
        if len(header + body) > room:
            truncated = True
        chunks.append(seg)
        used += len(seg)
    return "".join(chunks), truncated


def out_path() -> str:
    if ROLE == "adjudicator":
        return os.path.join(OUT_DIR, "ADJUDICATION.md")
    if ROLE == "extractor":
        return os.path.join(OUT_DIR, "extraction.json")
    return os.path.join(OUT_DIR, f"{GID}.json")


def write_out(text: str) -> None:
    os.makedirs(OUT_DIR, exist_ok=True)
    with open(out_path(), "w") as f:
        f.write(text.rstrip() + "\n")


def write_invalid_marker(err: str) -> None:
    try:
        os.makedirs(OUT_DIR, exist_ok=True)
        with open(out_path() + ".invalid", "w") as f:
            f.write(err.strip() + "\n")
    except Exception:
        pass


# ---------------------------------------------------------------- LLM client
def _post(payload: dict, timeout: int = 600) -> dict:
    req = urllib.request.Request(
        f"{LLM_BASE}/chat/completions",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {LLM_KEY}"},
        method="POST")
    with urllib.request.urlopen(req, timeout=timeout, context=_SSL_CTX) as resp:
        return json.loads(resp.read().decode())


def chat(messages: list, extra: dict | None = None) -> str:
    """One completion. Retries with backoff on 5xx/connection errors; raises
    HTTPError immediately on 4xx so the caller can walk the reliability ladder."""
    payload = {"model": MODEL, "messages": messages,
               "temperature": TEMPERATURE, "max_tokens": MAX_OUTPUT_TOKENS}
    if extra:
        payload.update(extra)
    last: Exception | None = None
    for attempt in range(1, TRANSPORT_RETRIES + 1):
        try:
            data = _post(payload)
            return (data["choices"][0]["message"].get("content") or "").strip()
        except urllib.error.HTTPError as e:
            if 400 <= e.code < 500:
                raise
            last = e
        except Exception as e:  # URLError, timeout, bad JSON envelope...
            last = e
        sleep = 5 * attempt
        log(f"transport error ({last}); retry {attempt}/{TRANSPORT_RETRIES} in {sleep}s")
        time.sleep(sleep)
    raise RuntimeError(f"model endpoint unreachable after {TRANSPORT_RETRIES} attempts: {last}")


def chat_json(messages: list, schema: dict) -> str:
    """The JSON reliability ladder: strict json_schema -> guided_json -> plain."""
    ladder = [
        ("json_schema", {"response_format": {"type": "json_schema", "json_schema": {
            "name": schema.get("title") or "output", "schema": schema, "strict": True}}}),
        ("guided_json", {"guided_json": schema}),
        ("plain", None),
    ]
    last: Exception | None = None
    for mode, extra in ladder:
        try:
            return chat(messages, extra=extra)
        except urllib.error.HTTPError as e:
            if 400 <= e.code < 500:
                log(f"{mode} request rejected (HTTP {e.code}); trying next mode")
                last = e
                continue
            raise
    raise RuntimeError(f"all completion modes were rejected by the endpoint: {last}")


# ---------------------------------------------------------------- JSON parse
def parse_json(text: str):
    """Whole-reply json.loads, else the first brace-balanced {...} block."""
    if not text:
        return None
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    start = text.find("{")
    while start != -1:
        depth, in_str, esc = 0, False, False
        for i in range(start, len(text)):
            c = text[i]
            if in_str:
                if esc:
                    esc = False
                elif c == "\\":
                    esc = True
                elif c == '"':
                    in_str = False
            elif c == '"':
                in_str = True
            elif c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(text[start:i + 1])
                    except json.JSONDecodeError:
                        break
        start = text.find("{", start + 1)
    return None


def structural_error(obj) -> str | None:
    """Lightweight in-agent checks; the full jsonschema pass is orchestrator-side."""
    if not isinstance(obj, dict):
        return "top-level value is not a JSON object"
    if ROLE == "extractor":
        for k in ("schema_version", "deck", "kpis", "forward_targets",
                  "red_flag_candidates", "narrative"):
            if k not in obj:
                return f"missing required key: {k}"
        if not isinstance(obj.get("deck"), dict) or "period" not in obj["deck"]:
            return "deck.period is missing"
    else:  # grader
        for k in ("schema_version", "grader", "categories", "red_flags", "overall_comment"):
            if k not in obj:
                return f"missing required key: {k}"
        cats = obj.get("categories")
        if not isinstance(cats, list) or len(cats) != 8:
            return "categories must contain exactly 8 entries (A-H)"
    return None


# ---------------------------------------------------------------- prompts
def _persona() -> str:
    return read_text(PERSONA_PATH).strip()


def _bdef() -> str:
    return read_text(BDEF_PATH).strip()


def _red_flag_taxonomy() -> str:
    """The '## Red-flag taxonomy' section of BDEF.md (falls back to the whole rubric)."""
    bdef = _bdef()
    low = bdef.lower()
    i = low.find("## red-flag taxonomy")
    return bdef[i:].strip() if i != -1 else bdef


def build_extractor_prompt(schema: dict) -> tuple[str, str, bool]:
    schema_text = json.dumps(schema, indent=2)
    system = (
        "You are the structured-data EXTRACTOR for a board-deck grading pipeline. "
        "You turn the deck text into machine-readable JSON; you do not grade.\n\n"
        "HARD RULES:\n"
        "- Reply with ONLY one JSON object conforming exactly to the schema below. "
        "No prose, no markdown fences, no comments.\n"
        "- NEVER invent numbers. Every actual/target must appear in the deck text; "
        "record where in \"source\".\n"
        "- canonical_name is lower_snake_case, generic, and stable across quarters "
        "(arr, ebitda_margin, churn_rate...).\n"
        "- deck.period is the reporting period as printed on the deck "
        "(e.g. 2026-Q2, 2026-H1, FY2026, 2026-05); null if truly absent.\n"
        "- If the deck text you were given was truncated, set deck.truncated = true.\n\n"
        "# OUTPUT SCHEMA (JSON Schema)\n" + schema_text
    )
    persona = _persona()
    if persona:
        system = persona + "\n\n" + system
    taxonomy = _red_flag_taxonomy()
    prefix = ("# RED-FLAG TAXONOMY\nUse ONLY these codes when populating "
              "red_flag_candidates:\n\n" + taxonomy + "\n\n# BOARD DECK TEXT\n")
    suffix = ("\n\n# TASK\nExtract the deck metadata, every KPI actual, every "
              "forward-looking target, red-flag candidates (taxonomy codes above), "
              "and the narrative summary. Output the JSON object now.")
    budget = max(8000, (MAX_MODEL_LEN - 3500) * 3 - len(system) - len(prefix) - len(suffix))
    docs, truncated = preload_docs(budget)
    note = ("\n\n(NOTE: the deck text above was TRUNCATED to fit the context window — "
            "set deck.truncated = true.)" if truncated else "")
    return system, prefix + docs + note + suffix, truncated


def build_grader_prompt(schema: dict) -> tuple[str, str, bool]:
    schema_text = json.dumps(schema, indent=2)
    system = (
        f"You are '{NAME}', one grader on a panel scoring a portfolio-company board "
        "deck against the BDEF rubric.\n\n"
        "HARD RULES:\n"
        "- Score every BDEF category A-H with an integer 1-5.\n"
        "- Any score ABOVE or BELOW 3 REQUIRES verbatim evidence quotes from the deck, "
        "each with a location (e.g. 'slide 6'). Unsupported non-3 scores will be "
        "regressed to 3 by the pipeline.\n"
        "- Do NOT compute totals or composite scores; numbers are computed elsewhere.\n"
        f"- Set \"grader\" to exactly \"{GID}\".\n"
        "- Use only the red-flag taxonomy codes defined in the rubric.\n"
        "- Reply with ONLY one JSON object conforming exactly to the provided schema. "
        "No prose, no markdown fences."
    )
    persona = _persona()
    if persona:
        system += "\n\n# YOUR LENS — how YOU specifically read this deck\n" + persona
    bdef = _bdef()
    prefix = "# BDEF RUBRIC\n" + bdef + "\n\n# BOARD DECK TEXT\n"
    suffix = ("\n\n# OUTPUT SCHEMA (JSON Schema)\n" + schema_text +
              "\n\n# TASK\nGrade the deck per the rubric and your lens. "
              "Output the JSON object now.")
    budget = max(8000, (MAX_MODEL_LEN - 3500) * 3 - len(system) - len(prefix) - len(suffix))
    docs, truncated = preload_docs(budget)
    note = ("\n\n(NOTE: the deck text above was TRUNCATED to fit the context window — "
            "grade what is shown and mention the truncation in overall_comment.)"
            if truncated else "")
    return system, prefix + docs + note + suffix, truncated


def _load_panel_grades() -> str:
    """All /grades/*.json panel reports, skipping extraction.json and anything
    flagged invalid by the agent that produced it (sibling .invalid marker)."""
    parts = []
    if not os.path.isdir(GRADES_DIR):
        return ""
    for fn in sorted(os.listdir(GRADES_DIR)):
        p = os.path.join(GRADES_DIR, fn)
        if not fn.endswith(".json") or not os.path.isfile(p):
            continue
        if fn == "extraction.json" or os.path.exists(p + ".invalid"):
            continue
        parts.append(f"\n\n========== GRADER REPORT: {fn} ==========\n" + read_text(p))
    return "".join(parts)


def build_adjudicator_prompt() -> tuple[str, str]:
    system = _persona() or (
        "You are the adjudicator chairing a panel of board-deck graders. You did not "
        "read the deck first-hand for a fresh opinion — you weigh the panel's evidence."
    )
    extraction_txt = read_text(EXTRACTION_PATH)
    grades_txt = _load_panel_grades()
    body = ("# STRUCTURED EXTRACTION (ground truth pulled from the deck)\n" +
            extraction_txt + "\n\n# PANEL GRADE REPORTS\n" + grades_txt)
    budget = max(8000, (MAX_MODEL_LEN - 3500) * 3 - len(system) - 1500)
    if len(body) > budget:
        body = body[:budget] + "\n\n(NOTE: input truncated to fit the context window.)"
    task = (
        "\n\n# TASK\nWrite a MARKDOWN adjudication of the panel:\n"
        "- Consensus per BDEF category A-H (one line each).\n"
        "- Material disagreements: where graders diverge, what each cites, and whose "
        "evidence is stronger (verbatim deck quotes beat assertions).\n"
        "- Red flags: which are CONFIRMED and which are DISMISSED, and why.\n"
        "- Exactly 3 questions the board should ask management next quarter.\n"
        "Do NOT output numeric scores, totals, or JSON — narrative Markdown only."
    )
    return system, body + task


# ---------------------------------------------------------------- run
def finalize(obj: dict, truncated: bool) -> dict:
    if ROLE == "grader":
        obj["grader"] = GID
    elif ROLE == "extractor" and truncated:
        deck = obj.get("deck")
        if isinstance(deck, dict):
            deck["truncated"] = True
    return obj


def run_json_role() -> None:
    schema = json.loads(read_text(SCHEMA_PATH) or "{}")
    if ROLE == "extractor":
        system, user, truncated = build_extractor_prompt(schema)
    else:
        system, user, truncated = build_grader_prompt(schema)
    if truncated:
        log("document text truncated to fit the context window")
    messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]

    text = chat_json(messages, schema)
    obj = parse_json(text)
    err = structural_error(obj) if obj is not None else "reply was not parseable JSON"
    if obj is not None and err is None:
        write_out(json.dumps(finalize(obj, truncated), indent=2))
        log(f"wrote {out_path()}")
        return

    # ONE repair round-trip.
    log(f"invalid reply ({err}); attempting one repair round-trip")
    repair = messages + [
        {"role": "assistant", "content": text or "(empty reply)"},
        {"role": "user", "content": (
            f"Your previous reply was not valid JSON or failed validation: {err}. "
            "Reply with ONLY the corrected JSON.")},
    ]
    text2 = chat_json(repair, schema)
    obj2 = parse_json(text2)
    err2 = structural_error(obj2) if obj2 is not None else "reply was not parseable JSON"
    if obj2 is not None and err2 is None:
        write_out(json.dumps(finalize(obj2, truncated), indent=2))
        log(f"wrote {out_path()} (after repair)")
        return

    # Still bad: leave the raw text + an .invalid marker for the orchestrator.
    write_out(text2 or text or "")
    write_invalid_marker(f"invalid after repair round-trip: {err2}")
    log(f"FAILED to produce valid JSON: {err2} (raw text + .invalid marker written)")


def run_adjudicator() -> None:
    system, user = build_adjudicator_prompt()
    messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
    text = chat(messages)
    if not text.strip():
        raise RuntimeError("model returned an empty adjudication")
    write_out(text)
    log(f"wrote {out_path()} ({len(text)} chars)")


def main() -> None:
    log(f"starting (role={ROLE} model={MODEL} base={LLM_BASE} temp={TEMPERATURE})")
    try:
        if ROLE == "adjudicator":
            run_adjudicator()
        else:
            run_json_role()
    except Exception as e:
        print(traceback.format_exc(), flush=True)
        if ROLE != "adjudicator":
            # Leave a marker so the orchestrator sees this agent ran and failed.
            write_invalid_marker(f"agent error: {e}")
        raise


if __name__ == "__main__":
    main()