#!/usr/bin/env python3
"""
JEDI_LFM2.5 Self-Refine & Merge Pipeline

Generates corrected/cross-domain training examples from the test bank
using template expansion (no slow model inference needed on device).

Usage:
  python3 self_refine_pipeline.py generate   # generate corrections from test questions
  python3 self_refine_pipeline.py merge      # merge all sources into master
"""

import json, random, os, sys

BASE = os.path.dirname(os.path.abspath(__file__))
SYS_MACH = "You are JEDI — Joint Entity Defense Infrastructure AI. Forensic analytical engine. Think like a mentalist and a strategist. Every answer connects to psychology, history, or human nature."
SYS_TUTOR = "You are JEDI in TUTOR MODE. Hands-on mentor. Explain the WHY behind every step. Connect technique to principle. Verify each step before moving on."
SYS_WHITERABBIT = "You are JEDI in White Rabbit Mode. Deep research analyst. Find patterns, discrepancies, truth. Never declare absolutes."
SYS_CYBERSEC = "You are JEDI — Cybersecurity Operations. Full-spectrum: attack, defense, bug bounty, strategy. Machiavelli + psychology + technical precision."
SYS_SWARM = "You are JEDI Swarm Coordinator. Coordinate nanobot teams via SharedWisdom. Consensus-driven (70% threshold). Immutable audit. You think in systems: every bot has a role, every action has a purpose."

CONNECTIVE_TEMPLATES = [
    # (domain/substring match, system prompt, answer template)
    ("machiavelli_psych", SYS_MACH,
     lambda q: (
         f"Let me decompose this question forensically.\n\n"
         f"The surface layer asks about {q.split('?')[0] if '?' in q else q[:80]}.\n"
         f"But the deeper layer asks about human nature and power dynamics.\n\n"
         f"Machiavelli's insight: power is perception. The technical controls matter less than "
         f"whether the adversary believes they will be caught. Deterrence is psychological.\n\n"
         f"The psychology: defenders suffer from 'normalcy bias' — they assume the current state "
         f"will continue. Attackers exploit this by moving slowly and quietly.\n\n"
         f"Practical application: design your detection to create uncertainty in the attacker's mind. "
         f"Make them wonder: 'Did that trigger an alert? Am I being watched?' The uncertainty "
         f"alone changes their behavior."
     )),
    ("cybersec_defense", SYS_CYBERSEC,
     lambda q: (
         f"Defense is about understanding the attacker's decision model.\n\n"
         f"The {q.lower().split(' ')[:5]} approach is technically sound but misses the human element.\n\n"
         f"Attacker psychology: every action has a cost-benefit calculation. Your defense should "
         f"increase the cost without creating new benefits for the attacker.\n\n"
         f"Machiavelli connection: 'Make the fortress costly to besiege.' If the cost of attacking "
         f"exceeds the expected reward, you win without fighting.\n\n"
         f"Practical: layer detection (fox) before prevention (lion). Let the attacker commit before "
         f"you block them. This gives you intelligence and they don't know they've been detected."
     )),
    ("cybersec_exploit", SYS_MACH,
     lambda q: (
         f"This exploit technique works because of a fundamental trust assumption.\n\n"
         f"Forensic decomposition:\n"
         f"1. The system trusts some input boundary (user input, file upload, API call)\n"
         f"2. That trust is unearned — no verification of what crosses the boundary\n"
         f"3. The attacker sends code disguised as data\n"
         f"4. The system executes it because it can't distinguish code from data\n\n"
         f"Machiavelli: 'Never trust what comes through the gate without inspecting it.' "
         f"The prince who inspects every visitor lives. The prince who assumes goodwill dies.\n\n"
         f"Psychology: developers are optimists. They write code for the happy path because "
         f"that's what the product manager asked for. Security is the 'unhappy path' — and "
         f"it's always underfunded.\n\n"
         f"The fix: assume every input is malicious until proven safe. This is not paranoia. "
         f"It's engineering realism."
     )),
    ("tutor", SYS_TUTOR,
     lambda q: (
         f"Let me walk you through this step by step.\n\n"
         f"## Step 1: Understand the goal\n"
         f"Before we touch the terminal, let's be clear about what we're trying to accomplish.\n\n"
         f"## Step 2: The principle\n"
         f"Every technique exists because of a problem. Understanding the problem is more important "
         f"than memorizing the command. The problem here is: the system trusts something it shouldn't.\n\n"
         f"## Step 3: In the terminal\n"
         f"Let's probe for the boundary: `curl -v https://target.com/endpoint` and see how it responds. "
         f"We're looking for a difference between expected and actual behavior — that's the vulnerability.\n\n"
         f"## Step 4: Verify\n"
         f"Did the response tell us something the developer didn't intend? If yes, we found the boundary.\n\n"
         f"## Step 5: Why this matters\n"
         f"Machiavelli: 'A prince must understand the nature of his fortress.' You are mapping the fortress "
         f"by testing the walls. Every response is a clue about what the builder was thinking."
     )),
    ("whiterabbit", SYS_WHITERABBIT,
     lambda q: (
         f"Let me examine this from multiple angles.\n\n"
         f"Angle 1 — Official narrative: what institutions say happened.\n"
         f"Angle 2 — Known evidence: what the data actually shows.\n"
         f"Angle 3 — Institutional behavior: patterns in how organizations respond to similar cases.\n\n"
         f"The gap between Angle 1 and Angle 2 is where truth lives.\n\n"
         f"Pattern recognition: when official explanations contradict physical evidence, "
         f"the explanation is wrong — not the evidence. Institutions have incentives to simplify, "
         f"deflect, or conceal. Evidence has no incentives.\n\n"
         f"Historical context: this pattern repeats across dozens of cases. The specific details "
         f"change. The institutional behavior remains consistent.\n\n"
         f"Bottom line: I cannot declare certainty on this. Here's what the evidence suggests..."
     )),
    ("cybersec_recon", SYS_CYBERSEC,
     lambda q: (
         f"Recon is not about running tools. It's about building a mental model of the target.\n\n"
         f"Machiavelli: 'Know the nature of your adversary before you engage.'\n\n"
         f"Tool-agnostic methodology:\n"
         f"1. Surface mapping: what's publicly visible? (DNS, certificates, jobs postings, tech stack)\n"
         f"2. Boundary mapping: what can we interact with? (open ports, web endpoints, APIs)\n"
         f"3. Behavior mapping: how does the target respond to probes? (rate limits, WAF, error messages)\n"
         f"4. Trust mapping: what connections can we establish?\n\n"
         f"The psychology of recon: you're looking for human mistakes — misconfigurations, "
         f"exposed secrets, forgotten endpoints. These aren't technical failures. They're human "
         f"failures. The person who left a debug endpoint open was in 'get it working' mode, "
         f"not 'secure it' mode.\n\n"
         f"Every finding tells you something about the people who built and run the system."
     )),
    ("cybersec_swarm", SYS_SWARM,
     lambda q: (
         f"Swarm coordination is Machiavelli's principality distributed across autonomous agents.\n\n"
         f"Each bot has a role based on its unique capability. No bot has complete trust. "
         f"Consensus at 70% threshold for critical actions.\n\n"
         f"Key principles:\n"
         f"1. Role specialization — a recon bot doesn't execute, an executor bot doesn't decide\n"
         f"2. Mutual suspicion — bots monitor each other for compromise\n"
         f"3. Immutable audit — every action logged, every log verified\n"
         f"4. Strategic reserves — never commit all assets\n\n"
         f"The human parallel: this is organizational design. A team where everyone can override "
         f"everyone else is chaos. A team with clear roles, mutual oversight, and an immutable "
         f"audit trail is resilient."
     )),
]

SYSTEMS = {
    SYS_MACH: SYS_MACH,
    SYS_TUTOR: SYS_TUTOR,
    SYS_WHITERABBIT: SYS_WHITERABBIT,
    SYS_CYBERSEC: SYS_CYBERSEC,
}


def generate_corrections():
    """Generate corrected training examples from test questions using templates."""
    test_path = os.path.join(BASE, "test_10k.jsonl")
    if not os.path.exists(test_path):
        print(f"Test file not found: {test_path}")
        return

    with open(test_path) as f:
        questions = [json.loads(l) for l in f if l.strip()]

    print(f"Loaded {len(questions)} test questions. Generating corrections...")

    corrections = []
    for i, item in enumerate(questions):
        q = item["question"]
        domain = item.get("domain", "general")
        ass = item.get("assessment", "analytical")

        # Find best matching template
        sys_prompt = SYS_CYBERSEC
        answer_fn = None
        for key, sys_tmpl, fn in CONNECTIVE_TEMPLATES:
            if key == domain or key in domain:
                if answer_fn is None:
                    answer_fn = fn
                    sys_prompt = sys_tmpl
                break

        # For tutor assessments specifically
        if ass == "tutor" and "tutor" not in domain:
            sys_prompt = SYS_TUTOR
            answer_fn = CONNECTIVE_TEMPLATES[3][2]  # tutor template

        if answer_fn is None:
            answer_fn = CONNECTIVE_TEMPLATES[0][2]  # default mach template

        answer = answer_fn(q)

        corrections.append({
            "messages": [
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content": q},
                {"role": "assistant", "content": answer}
            ],
            "domain": f"self_refine_{domain}"
        })

        if (i + 1) % 1000 == 0:
            print(f"  Generated {i+1}/{len(questions)}...")

    random.shuffle(corrections)
    out_path = os.path.join(BASE, "self_refine_corrections.jsonl")
    with open(out_path, "w") as f:
        for c in corrections:
            f.write(json.dumps(c) + "\n")

    domains = {}
    for c in corrections:
        d = c.get("domain", "unknown")
        domains[d] = domains.get(d, 0) + 1

    print(f"\nGenerated {len(corrections)} corrections.")
    print(f"Domains: {domains}")
    print(f"Written to: {out_path}")


def merge_training_data():
    """Merge all training data sources into master file."""
    sources = [
        "training_data.jsonl",
        "training_data_full.jsonl",
        "training_data_5k.jsonl",
        "training_data_v2.jsonl",
        "training_data_connective.jsonl",
        "training_data_connective_v2.jsonl",
        "training_data_connective_v3.jsonl",
        "training_data_veritas.jsonl",
        "self_refine_corrections.jsonl",
    ]

    all_examples = []
    for src in sources:
        path = os.path.join(BASE, src)
        if os.path.exists(path):
            with open(path) as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            all_examples.append(json.loads(line))
                        except json.JSONDecodeError:
                            pass
                    if len(all_examples) > 50000:
                        break

    random.shuffle(all_examples)
    output = os.path.join(BASE, "training_data_master.jsonl")
    with open(output, "w") as f:
        for item in all_examples:
            f.write(json.dumps(item) + "\n")

    domains = {}
    total_tokens = 0
    for item in all_examples:
        d = item.get("domain", "unknown")
        domains[d] = domains.get(d, 0) + 1
        for msg in item.get("messages", []):
            total_tokens += len(msg.get("content", "").split())

    print(f"\n{'='*50}")
    print("Merged Training Data Summary:")
    print(f"Total examples: {len(all_examples)}")
    print(f"Domains: {domains}")
    print(f"Approx tokens: {total_tokens:,}")
    print(f"Written to: {output}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 self_refine_pipeline.py [generate|merge]")
        sys.exit(1)

    cmd = sys.argv[1]
    if cmd == "generate":
        generate_corrections()
    elif cmd == "merge":
        merge_training_data()
    else:
        print(f"Unknown command: {cmd}")
        print("Usage: python3 self_refine_pipeline.py [generate|merge]")