import os
import sys
import json
import re
import torch
from datetime import datetime

# Add root directory to sys.path so we can import project modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data.bug_dataset import TRAINING_SCENARIOS
from orchestrator import Debugger
from env.codedebugger_env import CodeDebuggerEnv
from utils.prompts import get_simplified_prompt

# Unsloth & TRL imports
from unsloth import FastLanguageModel
from trl import GRPOConfig, GRPOTrainer
from datasets import Dataset

# 1. Setup Model via Unsloth (4-bit quantization for fast training)
max_seq_length = 2048
lora_rank = 16

print("Loading unsloth/Llama-3.2-1B-Instruct...")
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Llama-3.2-1B-Instruct",
    max_seq_length=max_seq_length,
    load_in_4bit=True,
    fast_inference=False, # Set to False since vLLM is not installed
)

# Apply LoRA Adapter
model = FastLanguageModel.get_peft_model(
    model,
    r=lora_rank,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha=lora_rank,
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

# 2. Setup Curriculum Dataset
# Train easy problems first, then medium, then hard
difficulty_map = {"easy": 0, "medium": 1, "hard": 2}
sorted_scenarios = sorted(TRAINING_SCENARIOS, key=lambda x: difficulty_map.get(x["difficulty"], 3))

def format_prompt(problem):
    sys_prompt = "Fix the python code. Return ONLY the fixed code in a markdown block."
    user_prompt = get_simplified_prompt(problem["buggy_code"], problem.get("error_type", "Bug"))
    return [
        {"role": "system", "content": sys_prompt},
        {"role": "user", "content": user_prompt}
    ]

dataset_dict = {
    "prompt": [format_prompt(p) for p in sorted_scenarios],
    "problem_id": [p["id"] for p in sorted_scenarios],
    "difficulty": [p["difficulty"] for p in sorted_scenarios]
}
train_dataset = Dataset.from_dict(dataset_dict)

# 3. Setup Rollout Reward Function
env = CodeDebuggerEnv(max_iterations=1)
global_step = 0

os.makedirs("outputs", exist_ok=True)
log_file = open("outputs/training_log.jsonl", "w")

def get_completion_text(comp):
    if isinstance(comp, list): return comp[0]["content"]
    if isinstance(comp, dict): return comp["content"]
    return str(comp)

def executor_reward(prompts, completions, problem_id, difficulty, **kwargs):
    global global_step
    rewards = []
    
    for i, comp in enumerate(completions):
        prob_id = problem_id[i][0] if isinstance(problem_id[i], list) else problem_id[i]
        diff = difficulty[i][0] if isinstance(difficulty[i], list) else difficulty[i]
        
        # Find corresponding problem
        prob = next(p for p in TRAINING_SCENARIOS if p["id"] == prob_id)
        
        # Extract code from LLM output
        completion_text = get_completion_text(comp)
        match = re.search(r"```(?:python)?\n?(.*?)\n?```", completion_text, re.DOTALL)
        fixed_code = match.group(1).strip() if match else completion_text.strip()
        
        # Execute code in env
        env.reset(prob)
        obs, reward_dict, done, info = env.step(fixed_code)
        
        total_reward = reward_dict["total"]
        rewards.append(total_reward)
        
        # Log to jsonl
        log_entry = {
            "step": global_step,
            "problem_id": prob_id,
            "difficulty": diff,
            "reward": total_reward,
            "test_score": reward_dict.get("test_score", 0.0),
            "total": total_reward
        }
        log_file.write(json.dumps(log_entry) + "\n")
        log_file.flush()
        
    global_step += 1
    return rewards

# 4. Training configuration
training_args = GRPOConfig(
    output_dir="outputs/grpo_training",
    learning_rate=5e-6,
    lr_scheduler_type="cosine",
    max_steps=100,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    optim="adamw_8bit",
    weight_decay=0.01,
    warmup_ratio=0.1,
    logging_steps=1,
    fp16=not torch.cuda.is_bf16_supported(),
    bf16=torch.cuda.is_bf16_supported(),
    report_to="none", # Disabling wandb for simplicity
)

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=executor_reward,
    args=training_args,
    train_dataset=train_dataset,
)

print("\nStarting GRPO Training...")
trainer.train()

# 5. CRITICAL: Save ONLY the adapter, do not merge
print("\nSaving LoRA adapter...")
model.save_pretrained("outputs/trained_adapter")
tokenizer.save_pretrained("outputs/trained_adapter")
print("Adapter saved successfully to outputs/trained_adapter")

# 6. Evaluate all 30 problems using the trained model
class LocalLLMFixer:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        FastLanguageModel.for_inference(self.model) # Enable native 2x faster inference
        
    def fix_code(self, buggy_code, error_type, description, test_cases, test_results=None, previous_explanation=None, iteration=1):
        prompt_text = get_simplified_prompt(buggy_code, error_type)
        messages = [
            {"role": "system", "content": "Fix the python code. Return ONLY the fixed code in a markdown block."},
            {"role": "user", "content": prompt_text}
        ]
        
        inputs = self.tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")
        outputs = self.model.generate(input_ids=inputs, max_new_tokens=1024, temperature=0.2)
        response = self.tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
        
        match = re.search(r"```(?:python)?\n?(.*?)\n?```", response, re.DOTALL)
        fixed_code = match.group(1).strip() if match else response.strip()
        
        return {"fixed_code": fixed_code, "method": "local_llm"}

print("\nEvaluating all 30 problems with trained model...")
debugger = Debugger(max_iterations=3)
debugger.fixer = LocalLLMFixer(model, tokenizer) # Hot-swap the fixer

trained_results = []
for i, prob in enumerate(TRAINING_SCENARIOS):
    print(f"[{i+1}/30] Evaluating {prob['title']}...")
    res = debugger.run(prob, verbose=False)
    trained_results.append(res)

output = {
    "format_version": 1,
    "model": "unsloth/Llama-3.2-1B-Instruct-GRPO",
    "timestamp": datetime.now().isoformat(),
    "results": trained_results,
}

with open("outputs/trained_scores.json", "w") as f:
    json.dump(output, f, indent=2)

# 7. Print summary comparison
try:
    with open("outputs/baseline_scores.json", "r") as f:
        b_data = json.load(f)
        baseline = b_data.get("results", b_data) if isinstance(b_data, dict) else b_data
        base_solved = sum(1 for p in baseline if p.get("solved"))
except Exception:
    base_solved = "?"

train_solved = sum(1 for r in trained_results if r["solved"])

print("\n" + "="*50)
print(f"TRAINING AND EVALUATION COMPLETE!")
print(f"Baseline Solved: {base_solved}/30")
print(f"Trained Solved:  {train_solved}/30")
print("="*50)
print("Dashboard ready: run `streamlit run app.py`")