"""Fine-tune DeepSeek R1 8B for the Exploit Agent using Unsloth + QLoRA.

Base: deepseek-ai/DeepSeek-R1-Distill-Llama-8B
Dataset: data/exploit/exploit_train.jsonl
Target: >85% valid AttackPlan JSON, expert score >4/5
Special: DO NOT filter <think> tokens — they are part of the training target.
"""

def train_exploit_model(
    base_model: str = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
    dataset_path: str = None,
    output_dir: str = "models/exploit-agent",
    epochs: int = 3,
    lora_r: int = 128,  # higher rank for complex reasoning
    lora_alpha: int = 16,
):
    print(f"=== Exploit Agent Fine-Tuning ===")
    print(f"Base: {base_model}")
    print(f"LoRA: r={lora_r} (higher for reasoning), alpha={lora_alpha}, epochs={epochs}")
    print(f"NOTE: <think> tokens are preserved in training data")

    try:
        from unsloth import FastLanguageModel
        from trl import SFTTrainer
        from transformers import TrainingArguments
        from datasets import load_dataset
    except ImportError:
        print("Install: pip install unsloth trl datasets transformers")
        return

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=base_model, max_seq_length=8192, load_in_4bit=True,
    )

    model = FastLanguageModel.get_peft_model(
        model, r=lora_r, lora_alpha=lora_alpha,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0, bias="none", use_gradient_checkpointing="unsloth",
    )

    import os
    if dataset_path is None:
        dataset_path = os.path.join(os.path.dirname(__file__), "..", "data", "exploit", "exploit_train.jsonl")

    dataset = load_dataset("json", data_files=dataset_path, split="train")

    def format_prompt(ex):
        # IMPORTANT: output includes <think> tags — do not strip them
        return f"### Instruction:\n{ex['instruction']}\n\n### Input:\n{ex['input']}\n\n### Response:\n{ex['output']}"

    trainer = SFTTrainer(
        model=model, tokenizer=tokenizer, train_dataset=dataset,
        args=TrainingArguments(
            output_dir=output_dir, num_train_epochs=epochs,
            per_device_train_batch_size=2, gradient_accumulation_steps=16,
            learning_rate=1e-4, warmup_steps=100, save_steps=500,
            fp16=True, report_to="none",
        ),
        formatting_func=format_prompt, max_seq_length=8192,
    )

    trainer.train()
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    print(f"Model saved to {output_dir}")


if __name__ == "__main__":
    train_exploit_model()