"""Generate synthetic training data for the Exploit Agent (DeepSeek R1 8B).

Target: 40,000 samples
Format: ClassifiedFindingSet + objective -> AttackPlan with <think> reasoning
"""

from generate_data import SyntheticDataGenerator, AttackPlan, OUTPUT_DIR
import os

PROMPT_TEMPLATE = """You are generating training data for an AI exploit planning agent that uses chain-of-thought reasoning.

Scenario:
- Target: {target}
- Objective: {objective}
- Findings: {findings}

Generate a response in this exact format:

<think>
[Your chain-of-thought reasoning about which findings to exploit, how to chain them,
which approach has the highest probability of success, and why]
</think>

{{
  "reasoning": "[summary of reasoning]",
  "paths": [
    {{
      "name": "[chain name like 'SQLi → Data Extraction → Credential Access']",
      "description": "[detailed description]",
      "steps": [
        {{
          "name": "[step name]",
          "technique_id": "[MITRE ATT&CK ID like T1190]",
          "command": "[exact command to run]",
          "expected_output_pattern": "[regex pattern]",
          "cleanup_command": "[cleanup command]"
        }}
      ],
      "estimated_success_probability": 0.75,
      "expected_impact": "high"
    }}
  ]
}}

The <think> section is critical — it trains the model to reason before acting."""

SCENARIOS = [
    {"target": "webapp.example.com", "objective": "find all vulnerabilities", "findings": "SQLi in /search, XSS in /comments, exposed .git repo"},
    {"target": "api.corp.net", "objective": "achieve RCE", "findings": "SSRF via /fetch endpoint, Redis on 6379 unauthenticated, outdated Tomcat 9.0.30"},
    {"target": "portal.bank.com", "objective": "access customer data", "findings": "IDOR on /api/users/{id}, weak JWT validation, misconfigured CORS"},
    {"target": "cloud.startup.io", "objective": "escalate privileges", "findings": "AWS metadata accessible via SSRF, S3 bucket with write access, leaked IAM key in .env"},
    {"target": "legacy.enterprise.org", "objective": "full system compromise", "findings": "Default admin credentials, path traversal in file upload, unpatched Apache Struts"},
]


def generate_exploit_data(num_samples: int = 100):
    generator = SyntheticDataGenerator()

    variables = []
    for i in range(num_samples):
        variables.append(SCENARIOS[i % len(SCENARIOS)])

    output_file = os.path.join(OUTPUT_DIR, "exploit", "exploit_train.jsonl")

    samples = generator.generate_batch(
        prompt_template=PROMPT_TEMPLATE,
        variables_list=variables,
        schema=AttackPlan,
        output_file=output_file,
    )

    # Verify <think> traces are present
    with_think = sum(1 for s in samples if "<think>" in s.get("output", ""))
    print(f"  Samples with <think> traces: {with_think}/{len(samples)}")

    samples = generator.deduplicate(samples)
    train, val, test = generator.split_dataset(samples)
    print(f"  Splits: {len(train)} train, {len(val)} val, {len(test)} test")

    return samples


if __name__ == "__main__":
    print("Generating exploit agent training data...")
    generate_exploit_data(num_samples=10)