import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments
)
from peft import LoraConfig
from trl import SFTTrainer

# --------------------------
# INTERNAL USE ONLY
# See PRIVATE_MODEL_TRAINING_NOTES.md for base model details
# --------------------------

# Configuration (internal use only)
OUTPUT_DIR = "./maeyen-ai-model"

# Use CPU (GPU incompatible)
device = "cpu"
print("Using CPU for training.")

# LoRA Configuration
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# Training Arguments
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    num_train_epochs=3,
    logging_steps=5,
    save_strategy="epoch",
    fp16=False,
    bf16=False,
    push_to_hub=False,
    report_to="none",
    use_cpu=True
)

# Load Model and Tokenizer
print("Loading model and tokenizer...")
# See PRIVATE_MODEL_TRAINING_NOTES.md for base model name
# Replace "BASE_MODEL_NAME_HERE" with actual base model from private notes
BASE_MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_NAME,
    device_map="cpu",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

# Synthetic Data
synthetic_data = [
    {
        "text": "<|im_start|>system\nYou are Maeyen AI Transaction Risk Agent. Assess risk and output valid JSON only with requires_human_review: true.<|im_end|>\n<|im_start|>user\nTransaction:\n- Seller verified: true\n- Buyer verified: false\n- Amount: 250000 NGN\n- Category: electronics\n- Seller transactions: 3\n- Seller dispute rate: 0.25\n- Evidence: product_photo, tracking_number\n- Missing: packing_video, serial_number<|im_end|>\n<|im_start|>assistant\n{\"risk_level\": \"high\", \"risk_score\": 82, \"reasons\": [\"High-value electronics transaction\", \"Seller has limited transaction history\", \"Seller dispute rate is high\", \"Serial number and packing video are missing\"], \"recommended_action\": \"Do not release payment. Request more delivery evidence and admin review.\", \"requires_human_review\": true}<|im_end|>"
    }
]
dataset = Dataset.from_list(synthetic_data)

# Train
print("Starting training...")
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    args=training_args,
    peft_config=lora_config,
    max_seq_length=1024
)

trainer.train()

# Save
trainer.model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print(f"Training complete! LoRA adapter saved to {OUTPUT_DIR}")
print("\nIMPORTANT: See PRIVATE_MODEL_TRAINING_NOTES.md for full details.")