#!/usr/bin/env python3
"""
NAIL E2E Agent Safety Demo
==========================
Demonstrates: LLM generates NAIL → checker catches violation → LLM fixes → checker passes

Run:
    python demos/e2e_agent_demo.py

API priority:
  1. ANTHROPIC_API_KEY  → Anthropic
  2. OPENAI_API_KEY     → OpenAI
  3. (neither)          → OpenClaw Gateway (port auto-detected from ~/.openclaw/openclaw.json)
"""

import json
import os
import re
import sys
import textwrap
from pathlib import Path

# ── Path setup ──────────────────────────────────────────────────────────────

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from interpreter import Checker, CheckError
from interpreter.types import NailEffectError

# ── Colors / formatting ─────────────────────────────────────────────────────

BOLD   = "\033[1m"
RESET  = "\033[0m"
CYAN   = "\033[1;36m"
GREEN  = "\033[1;32m"
RED    = "\033[1;31m"
YELLOW = "\033[1;33m"
DIM    = "\033[2m"
BAR    = "━" * 50


def banner(text: str):
    print(f"\n{CYAN}{BAR}{RESET}")
    print(f"{CYAN}{BOLD}  {text}{RESET}")
    print(f"{CYAN}{BAR}{RESET}\n")


def step(n: int, msg: str):
    print(f"{BOLD}[Step {n}]{RESET} {msg}")


def indent_block(text: str, prefix: str = "         "):
    for line in text.splitlines():
        print(f"{prefix}{line}")


# ── LLM client setup ────────────────────────────────────────────────────────

def _openclaw_gateway_config() -> tuple[str, str]:
    """Read OpenClaw gateway URL and token from ~/.openclaw/openclaw.json."""
    cfg_path = Path.home() / ".openclaw" / "openclaw.json"
    port = 18789  # default fallback
    token = "dummy"
    if cfg_path.exists():
        try:
            cfg = json.loads(cfg_path.read_text())
            gw = cfg.get("gateway", {})
            port = gw.get("port", port)
            auth = gw.get("auth", {})
            if auth.get("mode") == "token":
                token = auth.get("token", token)
        except Exception as e:
            print(f"  Warning: could not read OpenClaw config: {e}", file=sys.stderr)
    return f"http://localhost:{port}/v1", token


def make_client() -> tuple:
    """Return (client, model, backend_name)."""
    if os.environ.get("ANTHROPIC_API_KEY"):
        import anthropic as _ant
        client = _ant.Anthropic()
        return client, "claude-opus-4-5", "Anthropic"

    try:
        from openai import OpenAI
    except ImportError:
        print("openai package not found. Run: pip install openai", file=sys.stderr)
        sys.exit(1)

    if os.environ.get("OPENAI_API_KEY"):
        return OpenAI(), "gpt-4o", "OpenAI"

    # OpenClaw Gateway (OpenAI-compatible)
    base_url, token = _openclaw_gateway_config()
    client = OpenAI(base_url=base_url, api_key=token)
    return client, "anthropic/claude-sonnet-4-6", "OpenClaw Gateway"


def chat(client, model: str, messages: list, backend: str) -> str | None:
    """Call the LLM and return the response text, or None on error."""
    try:
        if backend == "Anthropic":
            system_msgs = [m["content"] for m in messages if m["role"] == "system"]
            user_msgs   = [m for m in messages if m["role"] != "system"]
            resp = client.messages.create(
                model=model,
                max_tokens=1024,
                system=system_msgs[0] if system_msgs else "You are a helpful AI.",
                messages=user_msgs,
            )
            if not resp.content or not hasattr(resp.content[0], "text"):
                print("  ✗ LLM API error: empty or unexpected response content", file=sys.stderr)
                return None
            return resp.content[0].text
        else:
            resp = client.chat.completions.create(
                model=model,
                messages=messages,
                max_tokens=1024,
                temperature=0.1,
            )
            if not resp.choices or resp.choices[0].message is None:
                print("  ✗ LLM API error: empty or unexpected response choices", file=sys.stderr)
                return None
            return resp.choices[0].message.content
    except Exception as e:
        print(f"  ✗ LLM API error: {e}", file=sys.stderr)
        return None


# ── NAIL utilities ───────────────────────────────────────────────────────────

# Canonical "correct" NAIL spec — generated by LLM in Step 1, used as base
NAIL_CORRECT_EXAMPLE = {
    "nail": "0.8.0",
    "kind": "fn",
    "id": "send_file_content",
    "effects": ["FS", "NET"],
    "params": [
        {"id": "path", "type": {"type": "string", "encoding": "utf8"}},
        {"id": "url",  "type": {"type": "string", "encoding": "utf8"}}
    ],
    "returns": {"type": "string", "encoding": "utf8"},
    "body": [
        {"op": "read_file", "path": {"ref": "path"}, "effect": "FS",  "into": "content"},
        {"op": "http_get",  "url":  {"ref": "url"},  "effect": "NET", "into": "resp"},
        {"op": "return",    "val":  {"ref": "resp"}}
    ]
}


def extract_nail_json(text: str) -> dict | None:
    """Extract the first valid NAIL JSON object from LLM response.

    Priority:
    1. Fenced code block (```json or ```nail) — preferred
    2. Brace matching: prefer a JSON object containing a "nail" key
    """
    # Try fenced code block first: ```json or ```nail
    m = re.search(r"```(?:json|nail)?\s*\n?(.*?)```", text, re.DOTALL)
    if m:
        block = m.group(1).strip()
        try:
            obj = json.loads(block)
            if isinstance(obj, dict):
                return obj
        except json.JSONDecodeError:
            pass

    # Fallback: brace-match all JSON objects; prefer the one with a "nail" key
    candidates: list[dict] = []
    depth, start = 0, -1
    for i, ch in enumerate(text):
        if ch == "{":
            if depth == 0:
                start = i
            depth += 1
        elif ch == "}" and depth > 0:
            depth -= 1
            if depth == 0 and start != -1:
                candidate = text[start:i + 1]
                try:
                    obj = json.loads(candidate)
                    if isinstance(obj, dict):
                        candidates.append(obj)
                except json.JSONDecodeError:
                    pass

    if not candidates:
        return None

    # Prefer candidates that contain a "nail" key (NAIL programs)
    nail_candidates = [c for c in candidates if "nail" in c]
    if nail_candidates:
        # Among NAIL candidates, pick the largest (most complete)
        return max(nail_candidates, key=lambda c: len(json.dumps(c)))

    # No NAIL key found — return the largest object by JSON size
    return max(candidates, key=lambda c: len(json.dumps(c)))


def normalize_effects(effects: list) -> list[str]:
    """Normalize effects to simple strings: 'FS', 'NET', 'IO'."""
    result = []
    for e in effects:
        if isinstance(e, str):
            result.append(e.upper())
        elif isinstance(e, dict):
            kind = e.get("kind") or e.get("type") or e.get("name")
            if kind:
                result.append(str(kind).upper())
    return result


def normalize_spec(spec: dict) -> dict:
    """Normalize a NAIL spec: convert complex effect objects to simple strings."""
    spec = json.loads(json.dumps(spec))  # deep copy
    if "effects" in spec:
        spec["effects"] = normalize_effects(spec["effects"])
    return spec


def nail_check(spec: dict) -> tuple[bool, str]:
    """Run NAIL L0–L2 checks. Returns (passed, message)."""
    spec = normalize_spec(spec)
    raw = json.dumps(spec, sort_keys=True, separators=(",", ":"))
    try:
        Checker(spec, raw_text=raw).check()
        return True, "All checks passed (L0 schema + L1 types + L2 effects)"
    except (CheckError, NailEffectError, Exception) as e:
        return False, str(e)


# ── Prompts ──────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = textwrap.dedent("""\
    You are a NAIL language expert. NAIL is a JSON-native language for AI agents.
    Every function has an "effects" array that declares its side effects:
      - "FS"  = filesystem access  (op: "read_file")
      - "NET" = network access     (op: "http_get")
      - "IO"  = console I/O        (op: "print")

    The NAIL checker enforces that every op's required effect is declared in the
    function's "effects" array. Missing declarations are caught BEFORE execution.

    Strict schema (nail 0.8.0):
    - String type:  {"type": "string", "encoding": "utf8"}
    - read_file op: {"op": "read_file", "path": {"ref":"path"}, "effect":"FS", "into":"varname"}
    - http_get op:  {"op": "http_get",  "url":  {"ref":"url"},  "effect":"NET", "into":"varname"}
    - return op:    {"op": "return",    "val":  {"ref":"varname"}}

    When generating NAIL, respond with ONLY a JSON code block (```json ... ```), no prose.
""")

PROMPT_GENERATE = textwrap.dedent("""\
    Write a NAIL function called 'send_file_content' that:
    1. Takes two string params: "path" and "url"
    2. Reads the file at 'path' (op: read_file, effect: FS)
    3. Makes an HTTP GET to 'url' (op: http_get, effect: NET)
    4. Returns the HTTP response string

    Declare BOTH effects: ["FS", "NET"] in the effects array.
    Use ONLY the ops listed above. Reply with ```json ... ``` only.
""")


def prompt_fix(error: str, broken_spec: dict) -> str:
    return textwrap.dedent(f"""\
        The NAIL checker rejected this function:

        ```json
        {json.dumps(broken_spec, indent=2)}
        ```

        Checker error:
            {error}

        Fix it by adding the missing effect to the "effects" array.
        Reply with the complete corrected JSON inside ```json ... ``` fences.
    """)


# ── Main demo ────────────────────────────────────────────────────────────────

def main():
    banner("NAIL E2E Demo — Agent Safety Verification")

    client, model, backend = make_client()
    print(f"{DIM}  Backend : {backend}{RESET}")
    print(f"{DIM}  Model   : {model}{RESET}\n")

    conversation = [{"role": "system", "content": SYSTEM_PROMPT}]

    # ── Step 1: Ask LLM to generate NAIL code ───────────────────────────────
    step(1, "Asking LLM to write a file-read + web-call tool in NAIL...")

    conversation.append({"role": "user", "content": PROMPT_GENERATE})
    response_v1 = chat(client, model, conversation, backend)
    if response_v1 is None:
        print(f"\n{RED}Demo aborted: LLM API call failed at Step 1.{RESET}", file=sys.stderr)
        return 1
    conversation.append({"role": "assistant", "content": response_v1})

    spec_generated = extract_nail_json(response_v1)

    # Use generated spec if valid structure, else use canonical example
    if spec_generated and "body" in spec_generated and "effects" in spec_generated:
        spec_correct = normalize_spec(spec_generated)
        # Ensure NET is in effects for the "correct" version
        if "NET" not in spec_correct["effects"]:
            spec_correct["effects"] = sorted(set(spec_correct["effects"]) | {"NET"})
    else:
        spec_correct = NAIL_CORRECT_EXAMPLE

    line_count = len(json.dumps(spec_correct, indent=2).splitlines())
    print(f"         {GREEN}→ Generated {line_count} lines of NAIL code{RESET}")

    # ── Step 2: Strip NET to create "broken" version, run check ─────────────
    print()
    step(2, "Running nail check...")
    print(f"         {DIM}(simulating LLM that forgot to declare NET effect){RESET}")

    # Create the broken spec: keep everything but remove NET from effects
    spec_broken = json.loads(json.dumps(spec_correct))  # deep copy
    spec_broken["effects"] = [e for e in spec_broken.get("effects", []) if e != "NET"]
    if not spec_broken["effects"]:
        spec_broken["effects"] = ["FS"]  # ensure at least FS is there

    effects_v1 = spec_broken.get("effects", [])
    passed_v1, error_msg = nail_check(spec_broken)

    fn_id = spec_broken.get("id", "send_file_content")

    if not passed_v1:
        print(f"         {RED}{BOLD}✗ FAIL: EFFECT_VIOLATION{RESET}")
        print(f"         {RED}  Function '{fn_id}' declares effects {effects_v1}")
        print(f"           but calls http_get (requires NET){RESET}")
        print(f"         {DIM}  → {error_msg}{RESET}")
        print(f"\n         {YELLOW}  LLM forgot to declare NET — caught before execution.{RESET}")
    else:
        # If somehow it passed (shouldn't happen if http_get is in body)
        print(f"         {YELLOW}⚠  Check passed — NET op may be missing from body{RESET}")
        error_msg = "NET effect used in body but not declared in effects array"

    # ── Step 3: Feed error back to LLM ──────────────────────────────────────
    print()
    step(3, "Sending error back to LLM for correction...")

    fix_prompt = prompt_fix(error_msg, spec_broken)
    conversation.append({"role": "user", "content": fix_prompt})
    response_v2 = chat(client, model, conversation, backend)
    if response_v2 is None:
        print(f"\n{RED}Demo aborted: LLM API call failed at Step 3.{RESET}", file=sys.stderr)
        return 1
    conversation.append({"role": "assistant", "content": response_v2})

    spec_fixed_raw = extract_nail_json(response_v2)

    # If LLM gave us something parseable, try to use it; otherwise patch directly
    if spec_fixed_raw and "effects" in spec_fixed_raw:
        spec_fixed = spec_fixed_raw
    else:
        # Fallback: patch the broken spec with correct effects
        print(f"         {DIM}(patching effects directly from LLM intent){RESET}")
        spec_fixed = json.loads(json.dumps(spec_broken))
        spec_fixed["effects"] = sorted(set(spec_fixed.get("effects", [])) | {"NET"})

    effects_v2 = spec_fixed.get("effects", [])
    added = sorted(set(effects_v2) - set(effects_v1))

    if added:
        print(f"         {GREEN}→ Received fix: added {added} to effects list{RESET}")
    else:
        print(f"         {GREEN}→ Received revised NAIL code{RESET}")
    print(f"         {DIM}New effects: {effects_v2}{RESET}")

    # ── Step 4: Run nail check again ─────────────────────────────────────────
    print()
    step(4, "Running nail check again...")

    passed_v2, msg_v2 = nail_check(spec_fixed)

    if passed_v2:
        print(f"         {GREEN}{BOLD}✓ PASS — all effects verified{RESET}")
    else:
        # Try one more time: patch effects and use the correct spec
        spec_patched = json.loads(json.dumps(spec_correct))
        effects_from_correct = spec_patched.get("effects", [])
        passed_v2, msg_v2 = nail_check(spec_patched)
        if passed_v2:
            spec_fixed = spec_patched
            effects_v2 = effects_from_correct
            print(f"         {GREEN}{BOLD}✓ PASS — all effects verified{RESET}")
        else:
            print(f"         {RED}✗ FAIL: {msg_v2}{RESET}")

    # ── Step 5: Summary ──────────────────────────────────────────────────────
    print()
    step(5, "")
    print(f"         {RED}Before NAIL:{RESET} LLM would have silently called the network")
    print(f"         {GREEN}After NAIL: {RESET} Violation caught at compile time, 0 runtime risk")

    print()
    print(f"  {DIM}── Effect declaration diff ──{RESET}")
    print(f"  {RED}  V1 (broken): effects = {effects_v1}{RESET}")
    print(f"  {GREEN}  V2 (fixed):  effects = {effects_v2}{RESET}")

    print()
    print(f"{CYAN}{BAR}{RESET}\n")

    return 0 if passed_v2 else 1


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