import argparse
import math
import os
import random
from typing import Any

import torch
import yaml
from peft import LoraConfig, PeftModel, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerBase
from universal_ml_utils.configuration import load_config
from universal_ml_utils.io import load_jsonl
from universal_ml_utils.logging import get_logger

from grasp.baselines.grisp.train import Lora
from grasp.configs import GraspConfig
from grasp.core import generate, load_notes, setup
from grasp.model import Message, Response, completions_api_messages
from grasp.rl.reward import RewardConfig, compute_reward, compute_step_rewards
from grasp.rl.vllm_engine import VllmEngine
from grasp.tasks import get_task


class OAPLConfig(GraspConfig):
    lora: Lora | None = None

    task: str = "cea"
    train_files: list[str]

    reward: RewardConfig = RewardConfig()

    beta1: float = 1.0
    beta2: float = 1e-3
    lr: float = 1e-5
    weight_decay: float = 0.01
    warmup_ratio: float = 0.05
    num_rollouts: int = 4
    num_epochs: int = 1
    batch_size: int = 4
    gradient_accumulation_steps: int = 1
    gradient_checkpointing: bool = False
    max_length: int = 8192

    # vLLM engine settings
    tensor_parallel_size: int = 1
    tool_call_parser: str = "qwen3_coder"
    use_ray: bool | None = None  # auto-detect: True if >1 GPU available
    vllm_kwargs: dict[str, Any] = {}

    log_interval: int = 10
    save_interval: int = 100


def gather_log_probs(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    # logits: (batch, seq, vocab), labels: (batch, seq)
    # shift: predict next token
    shift_logits = logits[:, :-1, :]
    shift_labels = labels[:, 1:]
    log_probs = torch.log_softmax(shift_logits, dim=-1)
    per_token = log_probs.gather(dim=-1, index=shift_labels.unsqueeze(-1)).squeeze(-1)
    return per_token


def estimate_v_star(rewards: list[float], beta1: float) -> float:
    G = len(rewards)
    if G == 0:
        return 0.0
    scaled = torch.tensor([r / beta1 for r in rewards])
    return beta1 * (torch.logsumexp(scaled, dim=0).item() - math.log(G))


def _reconstruct_messages(messages: list[dict]) -> list[Message]:
    msg_objects: list[Message] = []
    for msg in messages:
        role = msg["role"]
        content = msg["content"]
        if role == "assistant" and isinstance(content, dict):
            content = Response(**content)
        msg_objects.append(Message(role=role, content=content))
    return msg_objects


def _find_assistant_turn_spans(
    chat_messages: list[dict],
    tokenizer: PreTrainedTokenizerBase,
    seq_len: int,
) -> list[tuple[int, int]]:
    spans: list[tuple[int, int]] = []
    prefix_messages: list[dict] = []
    for msg in chat_messages:
        if msg["role"] == "assistant":
            prefix_len = _get_prefix_length(prefix_messages, tokenizer)
            prefix_messages.append(msg)
            full_len = _get_prefix_length(prefix_messages, tokenizer)
            start = min(prefix_len, seq_len)
            end = min(full_len, seq_len)
            if start < end:
                spans.append((start, end))
        else:
            prefix_messages.append(msg)
    return spans


def build_training_inputs(
    messages: list[dict],
    tokenizer: PreTrainedTokenizerBase,
    max_length: int,
    step_rewards: list[float] | None = None,
    v_star: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    msg_objects = _reconstruct_messages(messages)
    chat_messages = completions_api_messages(msg_objects)

    input_ids = tokenizer.apply_chat_template(
        chat_messages,
        return_tensors="pt",
        tokenize=True,
        add_generation_prompt=False,
    ).squeeze(0)

    if input_ids.shape[0] > max_length:
        input_ids = input_ids[:max_length]

    seq_len = input_ids.shape[0]

    # Find token spans for each assistant turn
    spans = _find_assistant_turn_spans(chat_messages, tokenizer, seq_len)

    # Build response mask and per-token advantages
    response_mask = torch.zeros(seq_len, dtype=torch.float32)
    token_advantages = torch.zeros(seq_len, dtype=torch.float32)

    if step_rewards is not None:
        # Process supervision: per-token advantage = return-to-go
        # (sum of current + future step rewards) minus V*
        cumulative_future = 0.0
        # Pair spans with step rewards, iterate in reverse for return-to-go
        num_steps = min(len(spans), len(step_rewards))
        returns_to_go = [0.0] * num_steps
        for i in range(num_steps - 1, -1, -1):
            cumulative_future += step_rewards[i]
            returns_to_go[i] = cumulative_future

        for i, (start, end) in enumerate(spans):
            response_mask[start:end] = 1.0
            if i < num_steps:
                token_advantages[start:end] = returns_to_go[i] - v_star
    else:
        # No step rewards: uniform advantage across all response tokens
        for start, end in spans:
            response_mask[start:end] = 1.0

    # Per-token ref logprobs: distribute vLLM logprobs across assistant spans
    ref_log_probs = torch.zeros(seq_len, dtype=torch.float32)
    assistant_responses = [
        msg_obj.content for msg_obj in msg_objects
        if msg_obj.role == "assistant" and isinstance(msg_obj.content, Response)
    ]
    for i, (start, end) in enumerate(spans):
        if i < len(assistant_responses) and assistant_responses[i].token_logprobs is not None:
            turn_logprobs = assistant_responses[i].token_logprobs
            span_len = end - start
            # vLLM logprobs may not align exactly token-for-token with HF;
            # pad or truncate to match span length
            if len(turn_logprobs) >= span_len:
                for j in range(span_len):
                    ref_log_probs[start + j] = turn_logprobs[j]
            else:
                for j, lp in enumerate(turn_logprobs):
                    if start + j < end:
                        ref_log_probs[start + j] = lp

    return input_ids, response_mask, ref_log_probs, token_advantages


def _get_prefix_length(messages: list[dict], tokenizer: PreTrainedTokenizerBase) -> int:
    if not messages:
        return 0
    ids = tokenizer.apply_chat_template(
        messages,
        return_tensors="pt",
        tokenize=True,
        add_generation_prompt=True,
    )
    return ids.shape[-1]


def sync_weights_to_vllm(
    train_model: PeftModel | torch.nn.Module,
    engine: VllmEngine,
) -> None:
    """Sync weights from training model to vLLM engine.

    Puts engine to sleep, broadcasts weights via NCCL, then wakes it up.
    """
    engine.sleep()
    engine.sync_weights(train_model)
    engine.wake_up()


def run_generate(
    task_name: str,
    input: object,
    grasp_config: GraspConfig,
    managers: list,
    kg_notes: dict[str, list[str]],
    notes: list[str],
    logger: object,
    engine: VllmEngine | None = None,
) -> dict:
    generator = generate(
        task_name,
        input,
        grasp_config,
        managers,
        kg_notes=kg_notes,
        notes=notes,
        logger=logger,  # type: ignore
        engine=engine,
    )

    output = None
    for output in generator:
        pass

    if output is None:
        raise RuntimeError("No output produced")

    return output


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="OAPL RL training for GRASP")
    parser.add_argument("config", type=str)
    parser.add_argument("output_dir", type=str)
    parser.add_argument("--log-level", type=str, default="INFO")
    return parser.parse_args()


def main(args: argparse.Namespace) -> None:
    logger = get_logger("OAPL", args.log_level)

    config = OAPLConfig(**load_config(args.config))

    assert config.seed is not None, "seed must be set for reproducible training"

    # Setup KG managers
    managers, _ = setup(config)
    notes, kg_notes = load_notes(config)

    # Load training data
    task_cls = get_task(config.task, managers, config)
    sample_cls = task_cls.__class__.sample_cls()
    assert sample_cls is not None, f"No sample class for task {config.task}"

    samples = []
    for f in config.train_files:
        for item in load_jsonl(f):
            samples.append(sample_cls(**item))

    logger.info(f"Loaded {len(samples)} training samples")

    # Resolve model name
    model_name = config.model

    # Load HF model for training
    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    if config.lora is not None:
        peft_config = LoraConfig(
            task_type="CAUSAL_LM",
            r=config.lora.r,
            lora_alpha=config.lora.lora_alpha,
            target_modules=config.lora.target_modules,
            lora_dropout=config.lora.dropout,
            modules_to_save=config.lora.save_modules,
        )
        model = get_peft_model(model, peft_config)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    model.train()

    if config.gradient_checkpointing:
        model.config.use_cache = False  # type: ignore
        if hasattr(model, "enable_input_require_grads"):
            model.enable_input_require_grads()
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})

    logger.info(f"Using model: {config.model} ({'LoRA' if config.lora else 'full'})")
    total = model.num_parameters()
    trainable = model.num_parameters(only_trainable=True)
    logger.info(f"Total: {total / 1e9:.1f}B, Trainable: {trainable / 1e6:.2f}M ({trainable / total:.2%})")

    # Create vLLM engine with dummy weights (real weights via NCCL)
    logger.info("Initializing vLLM engine...")
    engine = VllmEngine(
        model=model_name,
        tp_size=config.tensor_parallel_size,
        tool_call_parser=config.tool_call_parser,
        load_format="dummy",
        use_ray=config.use_ray,
        **config.vllm_kwargs,
    )
    logger.info(f"vLLM engine mode: {'Ray' if engine.use_ray else 'in-process'}")

    # Initialize NCCL weight transfer and broadcast initial weights
    logger.info("Initializing NCCL weight transfer...")
    engine.init_weight_transfer(model)
    sync_weights_to_vllm(model, engine)
    logger.info("Initial weight sync complete")

    # Optimizer
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=config.lr,
        weight_decay=config.weight_decay,
    )

    # Warmup + cosine schedule
    total_steps = math.ceil(
        len(samples) * config.num_epochs / config.batch_size / config.gradient_accumulation_steps
    )
    warmup_steps = int(total_steps * config.warmup_ratio)

    def lr_lambda(step: int) -> float:
        if step < warmup_steps:
            return step / max(1, warmup_steps)
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    # Wandb setup
    run_name = os.path.basename(args.output_dir)
    use_wandb = bool(os.environ.get("WANDB_PROJECT"))
    if use_wandb:
        import wandb

        os.environ["WANDB_NAME"] = run_name
        wandb.init(
            project=os.environ["WANDB_PROJECT"],
            name=run_name,
            config=config.model_dump(),
        )

    os.makedirs(args.output_dir, exist_ok=True)
    with open(os.path.join(args.output_dir, "config.yaml"), "w") as f:
        yaml.dump(config.model_dump(), f)

    global_step = 0
    random.seed(config.seed)
    torch.manual_seed(config.seed)

    for epoch in range(config.num_epochs):
        logger.info(f"Epoch {epoch + 1}/{config.num_epochs}")
        random.shuffle(samples)

        for batch_start in range(0, len(samples), config.batch_size):
            batch = samples[batch_start : batch_start + config.batch_size]

            # === Rollout phase: vLLM engine is active ===
            all_rollout_data: list[list[tuple[dict, float, list[float]]]] = []
            for sample in batch:
                sample_rollouts: list[tuple[dict, float, list[float]]] = []
                gt = [ann.model_dump() for ann in sample.annotations] if hasattr(sample, "annotations") else []

                for _ in range(config.num_rollouts):
                    try:
                        output = run_generate(
                            config.task,
                            sample.input(),
                            config,
                            managers,
                            kg_notes,
                            notes,
                            logger,
                            engine=engine,
                        )
                    except Exception as e:
                        logger.error(f"Rollout failed: {e}")
                        continue

                    reward = compute_reward(
                        output.get("output"),
                        gt,
                        output["messages"],
                        config.task,
                        config.reward,
                    )
                    step_rewards = compute_step_rewards(
                        output.get("output"),
                        gt,
                        output["messages"],
                        config.task,
                        config.reward,
                    )
                    sample_rollouts.append((output, reward, step_rewards))

                if sample_rollouts:
                    all_rollout_data.append(sample_rollouts)

            if not all_rollout_data:
                logger.warning("No successful rollouts in batch, skipping")
                continue

            # === Training phase ===
            num_total_rollouts = sum(len(sr) for sr in all_rollout_data)

            for sample_rollouts in all_rollout_data:
                rewards = [r for _, r, _ in sample_rollouts]
                v_star = estimate_v_star(rewards, config.beta1)

                for output, reward, step_rewards in sample_rollouts:
                    try:
                        input_ids, response_mask, ref_log_probs, token_advantages = build_training_inputs(
                            output["messages"], tokenizer, config.max_length,
                            step_rewards=step_rewards, v_star=v_star,
                        )
                    except Exception as e:
                        logger.error(f"Failed to build training inputs: {e}")
                        continue

                    input_ids = input_ids.to(device)
                    response_mask = response_mask.to(device)
                    ref_log_probs = ref_log_probs.to(device)
                    token_advantages = token_advantages.to(device)

                    logits = model(input_ids.unsqueeze(0)).logits
                    log_probs = gather_log_probs(logits, input_ids.unsqueeze(0)).squeeze(0)

                    # Align with shifted positions (gather_log_probs shifts by 1)
                    mask = response_mask[1:]
                    ref_lp = ref_log_probs[1:]
                    advantages = token_advantages[1:]

                    # Token-level OAPL: (β₂ * log_ratio_t - advantage_t)²
                    log_ratio = log_probs - ref_lp
                    residuals = config.beta2 * log_ratio - advantages
                    per_token_loss = residuals ** 2 * mask
                    loss = per_token_loss.sum() / (num_total_rollouts * config.gradient_accumulation_steps)
                    loss.backward()

            # Step optimizer every gradient_accumulation_steps
            if (global_step + 1) % config.gradient_accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
                scheduler.step()

            # === Sync phase: broadcast updated weights to vLLM via NCCL ===
            sync_weights_to_vllm(model, engine)

            global_step += 1

            # Logging
            if global_step % config.log_interval == 0:
                all_rewards = [r for sr in all_rollout_data for _, r, _ in sr]
                mean_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
                lr = scheduler.get_last_lr()[0]
                logger.info(
                    f"Step {global_step}/{total_steps} | "
                    f"reward={mean_reward:.4f} | lr={lr:.2e}"
                )
                if use_wandb:
                    wandb.log({
                        "reward/mean": mean_reward,
                        "reward/min": min(all_rewards) if all_rewards else 0.0,
                        "reward/max": max(all_rewards) if all_rewards else 0.0,
                        "lr": lr,
                        "global_step": global_step,
                    }, step=global_step)

            # Save checkpoint
            if global_step % config.save_interval == 0:
                ckpt_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                logger.info(f"Saving checkpoint to {ckpt_path}")
                model.save_pretrained(ckpt_path)
                tokenizer.save_pretrained(ckpt_path)

    # Final save
    final_path = os.path.join(args.output_dir, "final")
    logger.info(f"Saving final model to {final_path}")
    model.save_pretrained(final_path)
    tokenizer.save_pretrained(final_path)

    if use_wandb:
        wandb.finish()


if __name__ == "__main__":
    main(parse_args())
