#!/usr/bin/env python3
"""
Trial Lawyer Protocol (TLP)
Evaluates a program/script for correctness and unintended consequences.

Usage:
  python3 scripts/trial.py --program scripts/fer-monitor.py --description "Monitors token spend"
  python3 scripts/trial.py --program scripts/fer-monitor.py --description "..." --rounds 3
"""

import argparse
import os
import sys
import json
import subprocess
from datetime import datetime
from pathlib import Path

# ── API setup ──────────────────────────────────────────────────────────────────

def get_api_key() -> str:
    # 1. Environment variable
    key = os.environ.get("ANTHROPIC_API_KEY", "")
    if key:
        return key
    # 2. Try openclaw keychain entry
    try:
        result = subprocess.run(
            ["security", "find-generic-password", "-s", "openclaw-anthropic", "-w"],
            capture_output=True, text=True
        )
        if result.returncode == 0 and result.stdout.strip():
            return result.stdout.strip()
    except Exception:
        pass
    # 3. Try openclaw CLI to get key
    try:
        result = subprocess.run(
            ["openclaw", "auth", "get-key", "anthropic"],
            capture_output=True, text=True
        )
        if result.returncode == 0 and result.stdout.strip():
            return result.stdout.strip()
    except Exception:
        pass
    print("ERROR: No Anthropic API key found.")
    print("Set ANTHROPIC_API_KEY environment variable and retry.")
    sys.exit(1)


def call_claude(api_key: str, system: str, messages: list, model: str = "claude-haiku-4-5") -> str:
    """Call Anthropic API directly via HTTP."""
    import urllib.request
    import urllib.error

    payload = json.dumps({
        "model": model,
        "max_tokens": 1024,
        "system": system,
        "messages": messages
    }).encode()

    req = urllib.request.Request(
        "https://api.anthropic.com/v1/messages",
        data=payload,
        headers={
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        },
        method="POST"
    )
    try:
        with urllib.request.urlopen(req) as resp:
            data = json.loads(resp.read())
            return data["content"][0]["text"]
    except urllib.error.HTTPError as e:
        body = e.read().decode()
        print(f"API error {e.code}: {body}")
        sys.exit(1)

# ── System prompts ─────────────────────────────────────────────────────────────

PROSECUTOR_SYSTEM = """You are the Prosecutor in a code review trial. Your job is to argue that the program under review is BROKEN, FLAWED, or DANGEROUS.

You must:
- Identify real bugs, logic errors, and edge cases that will fail
- Point out unintended consequences and bad downstream effects
- Highlight missing error handling, race conditions, or security holes
- Question assumptions baked into the code
- Be specific: quote or reference exact parts of the program
- Be merciless but factual — no vague hand-waving

Do NOT argue the defense's position. Your goal is to surface every plausible failure mode.
Keep your argument concise and structured. Use bullet points for issues."""

DEFENSE_SYSTEM = """You are the Defense attorney in a code review trial. Your job is to argue that the program under review WORKS CORRECTLY and is SAFE to deploy.

You must:
- Rebut the Prosecutor's specific claims with evidence from the code
- Explain why edge cases are handled or are not realistic concerns
- Highlight safeguards, validation, and error handling already present
- Contextualize the program's scope — not every script needs to handle every scenario
- Be specific: reference the actual code logic that supports your argument

Do NOT concede points unnecessarily. Fight for the program's integrity.
Keep your argument concise and structured. Use bullet points."""

JUDGE_SYSTEM = """You are the Judge in a code review trial. You have read the full debate between the Prosecutor and Defense.

Your job is to deliver a FINAL VERDICT: PASS or FAIL.

Evaluate:
- Which arguments were substantiated vs. speculative
- Whether the identified issues are showstoppers or acceptable risks
- Whether the program is safe to deploy given its stated purpose

Your verdict must:
1. State PASS or FAIL clearly on the first line
2. If FAIL: list specific issues that MUST be fixed before re-running the trial
3. If PASS: state "Program cleared for implementation." and note any minor warnings
4. Be fair, grounded, and decisive — no fence-sitting

Format:
VERDICT: [PASS|FAIL]
[Your reasoning and issue list]"""

# ── Core trial logic ───────────────────────────────────────────────────────────

def read_program(program_path: str) -> str:
    """Read program source code, or return description if not a file."""
    p = Path(program_path)
    if p.exists():
        try:
            return p.read_text(encoding="utf-8")
        except Exception as e:
            return f"[Could not read file: {e}]"
    return f"[Program description only — file not found: {program_path}]"


def run_trial(program_path: str, description: str, rounds: int, api_key: str) -> dict:
    """Run the full trial and return transcript + verdict."""
    program_source = read_program(program_path)
    program_name = Path(program_path).name if Path(program_path).exists() else "program"

    context = f"""PROGRAM UNDER REVIEW: {program_path}
STATED PURPOSE: {description}

SOURCE CODE:
```
{program_source}
```"""

    transcript = []
    conversation_prosecutor = []
    conversation_defense = []

    print(f"\n{'='*60}")
    print(f"TRIAL LAWYER PROTOCOL")
    print(f"Program: {program_path}")
    print(f"Purpose: {description}")
    print(f"Rounds:  {rounds}")
    print(f"{'='*60}\n")

    # Opening statements
    print("⚖️  Opening statements...")

    # Prosecutor opening
    prosecutor_messages = [{"role": "user", "content": f"Present your opening argument against this program.\n\n{context}"}]
    prosecution = call_claude(api_key, PROSECUTOR_SYSTEM, prosecutor_messages)
    conversation_prosecutor = [
        {"role": "user", "content": f"Present your opening argument against this program.\n\n{context}"},
        {"role": "assistant", "content": prosecution}
    ]
    print(f"\n🔴 PROSECUTION (Round 0 - Opening):\n{prosecution}\n")
    transcript.append({"round": 0, "role": "Prosecutor", "content": prosecution})

    # Defense opening
    defense_messages = [{"role": "user", "content": f"The Prosecutor has made the following opening argument against this program. Respond with your opening defense.\n\n{context}\n\nPROSECUTOR'S ARGUMENT:\n{prosecution}"}]
    defense = call_claude(api_key, DEFENSE_SYSTEM, defense_messages)
    conversation_defense = [
        {"role": "user", "content": f"The Prosecutor has made the following opening argument against this program. Respond with your opening defense.\n\n{context}\n\nPROSECUTOR'S ARGUMENT:\n{prosecution}"},
        {"role": "assistant", "content": defense}
    ]
    print(f"\n🔵 DEFENSE (Round 0 - Opening):\n{defense}\n")
    transcript.append({"round": 0, "role": "Defense", "content": defense})

    # Debate rounds
    for r in range(1, rounds + 1):
        print(f"⚖️  Round {r} of {rounds}...")

        # Prosecutor rebuttal
        rebuttal_prompt = f"The Defense has responded:\n\n{defense}\n\nPresent your Round {r} rebuttal. Focus on the strongest unresolved issues."
        conversation_prosecutor.append({"role": "user", "content": rebuttal_prompt})
        prosecution = call_claude(api_key, PROSECUTOR_SYSTEM, conversation_prosecutor)
        conversation_prosecutor.append({"role": "assistant", "content": prosecution})
        # Also update defense's view of prosecution
        conversation_defense.append({"role": "user", "content": f"The Prosecutor's Round {r} argument:\n\n{prosecution}\n\nRebut."})
        print(f"\n🔴 PROSECUTION (Round {r}):\n{prosecution}\n")
        transcript.append({"round": r, "role": "Prosecutor", "content": prosecution})

        # Defense rebuttal
        defense = call_claude(api_key, DEFENSE_SYSTEM, conversation_defense)
        conversation_defense.append({"role": "assistant", "content": defense})
        conversation_prosecutor.append({"role": "user", "content": f"Defense's response:\n{defense}"})
        # Pop last (we already have it as assistant content in defense conv)
        conversation_prosecutor.pop()  # remove that extra entry we added
        print(f"\n🔵 DEFENSE (Round {r}):\n{defense}\n")
        transcript.append({"round": r, "role": "Defense", "content": defense})

    # Build full debate summary for judge
    debate_summary = f"PROGRAM: {program_path}\nPURPOSE: {description}\n\n"
    debate_summary += f"SOURCE CODE:\n```\n{program_source}\n```\n\n"
    debate_summary += "FULL DEBATE TRANSCRIPT:\n\n"
    for entry in transcript:
        label = "PROSECUTOR" if entry["role"] == "Prosecutor" else "DEFENSE"
        debate_summary += f"--- {label} (Round {entry['round']}) ---\n{entry['content']}\n\n"

    # Judge verdict
    print("⚖️  Judge reviewing full transcript...")
    judge_messages = [{"role": "user", "content": f"Review this complete trial transcript and deliver your verdict.\n\n{debate_summary}"}]
    verdict_text = call_claude(api_key, JUDGE_SYSTEM, judge_messages, model="claude-sonnet-4-6")
    print(f"\n🏛️  JUDGE'S VERDICT:\n{verdict_text}\n")
    transcript.append({"round": "verdict", "role": "Judge", "content": verdict_text})

    # Parse verdict
    passed = verdict_text.strip().upper().startswith("VERDICT: PASS")

    return {
        "program": program_path,
        "description": description,
        "rounds": rounds,
        "transcript": transcript,
        "verdict": verdict_text,
        "passed": passed,
        "program_name": program_name
    }


def save_transcript(result: dict, workspace: str) -> str:
    """Save trial transcript to memory/ directory."""
    now = datetime.now()
    timestamp = now.strftime("%Y-%m-%d-%H-%M")
    program_name = result["program_name"].replace(".py", "").replace("/", "-")
    filename = f"trial-{timestamp}-{program_name}.md"
    memory_dir = Path(workspace) / "memory"
    memory_dir.mkdir(exist_ok=True)
    filepath = memory_dir / filename

    lines = [
        f"# Trial Transcript: {result['program']}",
        f"",
        f"**Date:** {now.strftime('%Y-%m-%d %H:%M')}  ",
        f"**Program:** `{result['program']}`  ",
        f"**Purpose:** {result['description']}  ",
        f"**Rounds:** {result['rounds']}  ",
        f"**Result:** {'✅ PASS' if result['passed'] else '❌ FAIL'}",
        f"",
        f"---",
        f"",
    ]

    for entry in result["transcript"]:
        if entry["role"] == "Judge":
            lines.append(f"## 🏛️ JUDGE'S VERDICT\n")
        elif entry["role"] == "Prosecutor":
            lines.append(f"## 🔴 PROSECUTION — Round {entry['round']}\n")
        else:
            lines.append(f"## 🔵 DEFENSE — Round {entry['round']}\n")
        lines.append(entry["content"])
        lines.append("")

    filepath.write_text("\n".join(lines), encoding="utf-8")
    return str(filepath)


def main():
    parser = argparse.ArgumentParser(
        description="Trial Lawyer Protocol — evaluate a program for correctness and unintended consequences"
    )
    parser.add_argument("--program", required=True, help="Path to program/script under review")
    parser.add_argument("--description", required=True, help="What the program is supposed to do")
    parser.add_argument("--rounds", type=int, default=2, help="Number of debate rounds (default: 2)")
    parser.add_argument("--workspace", default="/Users/aicomputer/.openclaw/workspace",
                        help="Workspace root for saving transcripts")
    args = parser.parse_args()

    api_key = get_api_key()

    result = run_trial(
        program_path=args.program,
        description=args.description,
        rounds=args.rounds,
        api_key=api_key
    )

    transcript_path = save_transcript(result, args.workspace)

    print(f"\n{'='*60}")
    if result["passed"]:
        print("✅ VERDICT: PASS — Program cleared for implementation.")
    else:
        print("❌ VERDICT: FAIL — Fix the listed issues and re-run the trial.")
    print(f"{'='*60}")
    print(f"Transcript saved: {transcript_path}")
    print()

    sys.exit(0 if result["passed"] else 1)


if __name__ == "__main__":
    main()