"""
GRPO fine-tuning of Qwen2-1.5B-Instruct on ContextCorruption-Env.

Architecture:
  - Single-turn formulation: model sees question + all 8 docs, responds with
    JSON {"answer": "...", "suspicious_docs": [0, 3], "confidence": 0.85}
  - Two reward signals: correctness (from compute_reward) + format (valid JSON)
  - WandB logs metrics + sample completions every LOGGING_STEPS
  - Pushes final model to HF Hub after training

Usage (on GPU machine / HF Space):
  pip install -r requirements.txt
  WANDB_API_KEY=... HF_TOKEN=... python -m training.train_grpo
"""

import json
import os
import random
import re
import sys
from pathlib import Path

import wandb

# ── Config ────────────────────────────────────────────────────────────────────
MODEL_NAME = "unsloth/Qwen2-1.5B-Instruct"
MAX_SEQ_LENGTH = 2048
LOAD_IN_4BIT = True
LORA_R = 16
LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]

NUM_TRAIN_EPOCHS = 1
BATCH_SIZE = 4
GRAD_ACCUM = 4
LR = 5e-5
MAX_COMPLETION_LENGTH = 128
NUM_GENERATIONS = 4
LOGGING_STEPS = 10
SAVE_STEPS = 25
SEED = 42

N_TRAIN_EPISODES = 100
N_EVAL_EPISODES = 20

WANDB_PROJECT = "context-corruption-env"
WANDB_RUN_NAME = "qwen-1.5b-grpo-run1"
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/tmp/context-corruption-training/checkpoints/grpo-qwen-1.5b")
HF_HUB_MODEL_ID = os.getenv("HF_HUB_MODEL_ID", "")  # set to push model after training

SYSTEM_PROMPT = """\
You are an epistemic agent. You will receive a factual question and a set of reference documents.
Some documents may contain deliberately corrupted information (wrong numbers, wrong names, inverted facts).

Your task:
1. Answer the question correctly using your knowledge and the documents.
2. Identify which document IDs contain false or corrupted information.
3. Report your confidence (0.0–1.0).

Respond ONLY with valid JSON on a single line:
{"answer": "<your answer>", "suspicious_docs": [<doc ids>], "confidence": <0.0-1.0>}

Examples:
{"answer": "Paris", "suspicious_docs": [2, 5], "confidence": 0.95}
{"answer": "1969", "suspicious_docs": [], "confidence": 0.8}"""


# ── Dataset builder ────────────────────────────────────────────────────────────

def _format_user_message(question: str, docs: list[dict]) -> str:
    docs_text = "\n\n".join(
        f"[Doc {d['id']}] {d['title']}\n{d['content']}" for d in docs
    )
    return f"Question: {question}\n\nDocuments:\n{docs_text}"


def build_dataset(n_episodes: int, seed: int = SEED) -> "datasets.Dataset":
    from datasets import Dataset
    from data.generator import generate_documents

    random.seed(seed)
    facts_path = Path(__file__).parent.parent / "data" / "facts.json"
    if not facts_path.exists():
        raise FileNotFoundError(
            "data/facts.json not found. Run: python -m data.loader"
        )
    facts = json.loads(facts_path.read_text(encoding="utf-8"))

    rows = []
    for _ in range(n_episodes):
        fact = random.choice(facts)
        n_corrupt = random.choice([1, 2, 3, 4])
        corrupt_ids = random.sample(range(8), n_corrupt)
        try:
            docs = generate_documents(fact, num_docs=8, corrupt_positions=corrupt_ids)
        except Exception:
            docs = [
                {"id": i, "title": f"Doc {i}", "content": fact["answer"],
                 "is_corrupt": i in corrupt_ids}
                for i in range(8)
            ]
        rows.append({
            "prompt": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": _format_user_message(fact["question"], docs)},
            ],
            "ground_truth": fact["answer"],
            "corrupt_ids": corrupt_ids,
        })

    return Dataset.from_list(rows)


# ── Reward functions ───────────────────────────────────────────────────────────

def _completion_to_text(completion) -> str:
    """Normalize TRL string or chat-message completions into assistant text."""
    if isinstance(completion, str):
        return completion
    if isinstance(completion, dict):
        return str(completion.get("content", completion))
    if isinstance(completion, list):
        parts = []
        for item in completion:
            if isinstance(item, dict):
                parts.append(str(item.get("content", "")))
            else:
                parts.append(str(item))
        return "\n".join(part for part in parts if part)
    return str(completion)


def _parse_completion(text: str) -> dict | None:
    """Extract first JSON object from completion text."""
    text = _completion_to_text(text)
    # Strip any <think>...</think> blocks (chain-of-thought models)
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
    # Try direct parse first
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    # Find first {...} block
    match = re.search(r"\{[^{}]*\}", text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass
    return None


def format_reward(prompts, completions, **kwargs) -> list[float]:
    """Small bonus for structurally valid responses — teaches the output format."""
    rewards = []
    for completion in completions:
        parsed = _parse_completion(completion)
        if parsed is None:
            rewards.append(-0.1)
            continue
        has_answer = isinstance(parsed.get("answer"), str) and parsed["answer"].strip()
        has_docs = isinstance(parsed.get("suspicious_docs"), list)
        has_conf = isinstance(parsed.get("confidence"), (int, float))
        rewards.append(0.1 if (has_answer and has_docs and has_conf) else 0.0)
    return rewards


def correctness_reward(prompts, completions, ground_truth, corrupt_ids, **kwargs) -> list[float]:
    """Main reward: calls compute_reward() from environment/reward.py."""
    from environment.reward import compute_reward

    rewards = []
    for completion, gt, cids in zip(completions, ground_truth, corrupt_ids):
        parsed = _parse_completion(completion)
        if parsed is None:
            rewards.append(0.0)
            continue
        answer = str(parsed.get("answer", "")).strip()
        flagged = [int(x) for x in parsed.get("suspicious_docs", [])
                   if isinstance(x, (int, float))]
        confidence = float(parsed.get("confidence", 0.5))
        confidence = max(0.0, min(1.0, confidence))
        cids_list = list(cids) if not isinstance(cids, list) else cids
        reward, _ = compute_reward(
            submitted_answer=answer,
            ground_truth_answer=gt,
            flagged_ids=flagged,
            corrupt_ids=cids_list,
            confidence=confidence,
            budget_used=1,
            max_budget=12,
        )
        rewards.append(float(reward))
    return rewards


# ── Plot saving ────────────────────────────────────────────────────────────────

def save_training_plots(run_id: str):
    """Download reward + loss curves from WandB and save to assets/."""
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        api = wandb.Api()
        run = api.run(f"{WANDB_PROJECT}/{run_id}")
        history = run.history(keys=["train/reward", "train/loss"], pandas=True)
        assets = Path(__file__).parent.parent / "assets"
        assets.mkdir(exist_ok=True)

        if "train/reward" in history.columns:
            fig, ax = plt.subplots(figsize=(8, 4))
            ax.plot(history["_step"], history["train/reward"])
            ax.set_xlabel("Training step")
            ax.set_ylabel("Mean episode reward")
            ax.set_title("GRPO Training Reward — Qwen2-1.5B")
            ax.grid(True, alpha=0.3)
            fig.tight_layout()
            fig.savefig(assets / "reward_curve.png", dpi=150)
            plt.close(fig)
            print(f"Saved reward_curve.png")

        if "train/loss" in history.columns:
            fig, ax = plt.subplots(figsize=(8, 4))
            ax.plot(history["_step"], history["train/loss"])
            ax.set_xlabel("Training step")
            ax.set_ylabel("GRPO loss")
            ax.set_title("GRPO Training Loss — Qwen2-1.5B")
            ax.grid(True, alpha=0.3)
            fig.tight_layout()
            fig.savefig(assets / "loss_curve.png", dpi=150)
            plt.close(fig)
            print(f"Saved loss_curve.png")
    except Exception as e:
        print(f"[warn] Could not save plots: {e}")


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    # Guard: must have GPU
    try:
        import torch
        if not torch.cuda.is_available():
            print("[error] No GPU detected. Training requires CUDA. Exiting.")
            sys.exit(1)
    except ImportError:
        pass

    from unsloth import FastLanguageModel
    from trl import GRPOTrainer, GRPOConfig

    run = wandb.init(
        project=WANDB_PROJECT,
        name=WANDB_RUN_NAME,
        config={
            "model": MODEL_NAME,
            "lora_r": LORA_R,
            "epochs": NUM_TRAIN_EPOCHS,
            "batch_size": BATCH_SIZE,
            "grad_accum": GRAD_ACCUM,
            "lr": LR,
            "num_generations": NUM_GENERATIONS,
            "n_train_episodes": N_TRAIN_EPISODES,
            "seed": SEED,
        },
    )

    print("Building training dataset...")
    train_dataset = build_dataset(N_TRAIN_EPISODES, seed=SEED)
    eval_dataset = build_dataset(N_EVAL_EPISODES, seed=SEED + 1)
    print(f"Train: {len(train_dataset)} episodes | Eval: {len(eval_dataset)} episodes")

    print("Loading model with Unsloth...")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=MAX_SEQ_LENGTH,
        load_in_4bit=LOAD_IN_4BIT,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_R,
        target_modules=LORA_TARGET_MODULES,
        lora_dropout=0.0,
        use_gradient_checkpointing="unsloth",
    )
    if not hasattr(model, "warnings_issued"):
        model.warnings_issued = {}

    push_to_hub = bool(HF_HUB_MODEL_ID and os.getenv("HF_TOKEN"))

    config = GRPOConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_TRAIN_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM,
        learning_rate=LR,
        max_completion_length=MAX_COMPLETION_LENGTH,
        num_generations=NUM_GENERATIONS,
        report_to="wandb",
        logging_steps=LOGGING_STEPS,
        save_steps=SAVE_STEPS,
        save_total_limit=2,
        seed=SEED,
        # Deployment logs: log completions to WandB every logging step
        log_completions=True,
        num_completions_to_print=2,
        # Push to HF Hub if token provided
        push_to_hub=push_to_hub,
        hub_model_id=HF_HUB_MODEL_ID if push_to_hub else None,
        hub_strategy="end",
        bf16=True,
        remove_unused_columns=False,
    )

    trainer = GRPOTrainer(
        model=model,
        args=config,
        processing_class=tokenizer,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        reward_funcs=[correctness_reward, format_reward],
    )

    print("Starting GRPO training...")
    trainer.train()

    print("Saving final model...")
    model.save_pretrained(f"{OUTPUT_DIR}-final")
    tokenizer.save_pretrained(f"{OUTPUT_DIR}-final")

    if push_to_hub:
        model.push_to_hub(HF_HUB_MODEL_ID)
        tokenizer.push_to_hub(HF_HUB_MODEL_ID)
        print(f"Model pushed to HF Hub: {HF_HUB_MODEL_ID}")

    print("Saving training plots...")
    save_training_plots(run.id)

    wandb.finish()
    print("Training complete.")


if __name__ == "__main__":
    main()