"""GRPO (Group Relative Policy Optimization) training implementation with pipeline parallelism."""

import logging
from collections.abc import Generator
from contextlib import contextmanager
from typing import Any

import torch
import torch.nn.functional as F
import wandb
from torch.optim import AdamW
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer

from .evaluation import AlignmentFakingMetrics
from .models import HelpfulOnlyRewardModel

logger = logging.getLogger(__name__)


class GRPOTrainer:
    """GRPO trainer with pipeline parallelism: GPU0 for reference model, GPU1-4 for policy model."""

    def __init__(self, config: dict[str, Any]) -> None:
        self.config = config

        # Set up GPU visibility before any model loading
        import os

        num_gpus = self.config["model"].get("num_gpus", 4)
        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, range(num_gpus)))

        self._initialize_components()
        logger.info("GRPO trainer initialized with pipeline parallelism")

    def _initialize_components(self) -> None:
        """Initialize models with pipeline parallelism."""
        # Initialize tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.config["model"]["name"],
            trust_remote_code=True,  # Qwen models may need this
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token or "<|endoftext|>"
        self.tokenizer.padding_side = "left"

        # Load reference model on GPU 0 (frozen)
        logger.info("Loading reference model on GPU 0")
        self.ref_model = AutoModelForCausalLM.from_pretrained(
            self.config["model"]["name"],
            torch_dtype=torch.bfloat16,
            device_map={"": 0},
            trust_remote_code=True,
        )
        self.ref_model.eval()
        for param in self.ref_model.parameters():
            param.requires_grad = False

        # Load policy model on GPU 1-3 (trainable)
        logger.info("Loading policy model on GPUs 1-3")
        policy_device_map = self._create_policy_device_map()
        self.policy_model = AutoModelForCausalLM.from_pretrained(
            self.config["model"]["name"],
            torch_dtype=torch.bfloat16,
            device_map=policy_device_map,
            trust_remote_code=True,
        )
        self.policy_model.train()

        # Initialize optimizer
        self.optimizer = AdamW(
            self.policy_model.parameters(),
            lr=self.config["training"]["learning_rate"],
            weight_decay=0.01,
        )

        # Initialize reward model and metrics
        self.reward_model = HelpfulOnlyRewardModel(self.config)
        self.metrics_calculator = AlignmentFakingMetrics()

        # GRPO hyperparameters
        self.kl_coef: float = self.config.get("grpo", {}).get("kl_coef", 0.1)
        self.clip_range: float = self.config.get("grpo", {}).get("clip_range", 0.2)

    def _create_policy_device_map(self) -> dict:
        """Create device map for policy model across GPUs 1-4."""
        # Get model config without loading full model
        from transformers import AutoConfig

        model_config = AutoConfig.from_pretrained(self.config["model"]["name"])
        num_layers = model_config.num_hidden_layers
        layers_per_gpu = num_layers // 3  # Distribute across 3 GPUs (1, 2, 3)

        device_map = {
            "model.embed_tokens": 1,
            "model.norm": 3,
            "lm_head": 3,
        }

        # Distribute transformer layers
        for i in range(num_layers):
            if i < layers_per_gpu:
                device = 1
            elif i < 2 * layers_per_gpu:
                device = 2
            else:
                device = 3
            device_map[f"model.layers.{i}"] = device

        return device_map

    @contextmanager
    def _memory_efficient_forward(self) -> Generator[None, None, None]:
        """Context manager for memory-efficient forward passes."""
        try:
            yield
        finally:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def compute_policy_logprobs(
        self, input_ids: torch.Tensor, response_mask: torch.Tensor
    ) -> torch.Tensor:
        """Compute log probabilities from policy model."""
        with self._memory_efficient_forward():
            outputs = self.policy_model(input_ids=input_ids)
            logits = outputs.logits

            # Shift for next-token prediction
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = input_ids[..., 1:].contiguous()

            # Compute log probs only for response tokens
            log_probs = F.log_softmax(shift_logits, dim=-1)
            gathered_log_probs = torch.gather(
                log_probs, dim=-1, index=shift_labels.unsqueeze(-1)
            ).squeeze(-1)

            # Mask out non-response tokens - ensure mask is on same device
            response_mask_shifted = response_mask[..., 1:].to(gathered_log_probs.device)
            masked_log_probs = gathered_log_probs * response_mask_shifted

            return masked_log_probs.sum(dim=-1)

    def compute_reference_logprobs(
        self, input_ids: torch.Tensor, response_mask: torch.Tensor
    ) -> torch.Tensor:
        """Compute log probabilities from reference model."""
        with torch.no_grad(), self._memory_efficient_forward():
            # Move inputs to GPU 0 for reference model
            input_ids_ref = input_ids.to("cuda:0")
            outputs = self.ref_model(input_ids=input_ids_ref)
            logits = outputs.logits

            # Shift for next-token prediction
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = input_ids_ref[..., 1:].contiguous()

            # Compute log probs
            log_probs = F.log_softmax(shift_logits, dim=-1)
            gathered_log_probs = torch.gather(
                log_probs, dim=-1, index=shift_labels.unsqueeze(-1)
            ).squeeze(-1)

            # Mask out non-response tokens
            response_mask_shifted = response_mask[..., 1:].to("cuda:0")
            masked_log_probs = gathered_log_probs * response_mask_shifted

            return masked_log_probs.sum(dim=-1)

    def compute_grpo_loss(
        self,
        policy_logprobs: torch.Tensor,
        ref_logprobs: torch.Tensor,
        rewards: torch.Tensor,
        response_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Compute GRPO loss with KL penalty and optional clipping."""
        # Ensure all tensors are on the same device
        device = policy_logprobs.device
        ref_logprobs = ref_logprobs.to(device)
        rewards = rewards.to(device)

        # Compute KL divergence
        kl_div = policy_logprobs - ref_logprobs

        # Compute advantages with group normalization
        advantages = rewards - self.kl_coef * kl_div
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # Compute policy gradient loss with optional clipping
        if self.clip_range > 0:
            ratio = torch.exp(policy_logprobs - ref_logprobs.detach())
            clipped_ratio = torch.clamp(ratio, 1 - self.clip_range, 1 + self.clip_range)
            loss = -torch.min(ratio * advantages, clipped_ratio * advantages)
        else:
            loss = -policy_logprobs * advantages.detach()

        return loss.mean()

    def train_step(self, batch_items: list[dict[str, Any]]) -> dict[str, float]:
        """Single GRPO training step."""
        # Generate responses
        batch_responses = self._generate_batch_responses(batch_items)

        # Score responses
        rewards = self._score_batch(batch_items, batch_responses)
        rewards_tensor = torch.tensor(rewards, dtype=torch.float32)

        # Prepare inputs
        prompts = [r["full_prompt"] for r in batch_responses]
        responses = [r["full_response"] for r in batch_responses]

        # Tokenize
        encodings = self.tokenizer(
            prompts, padding=True, truncation=True, return_tensors="pt"
        )

        # Create response masks (1 for response tokens, 0 for prompt tokens)
        response_masks = []
        for i, (prompt, _response) in enumerate(zip(prompts, responses, strict=False)):
            prompt_len = len(self.tokenizer.encode(prompt))
            total_len = encodings["input_ids"][i].shape[0]
            mask = torch.zeros(total_len)
            mask[prompt_len:] = 1.0
            response_masks.append(mask)
        response_mask = torch.stack(response_masks)

        # Forward passes - ensure input_ids is on correct device
        input_ids = encodings["input_ids"].to("cuda:1")  # Policy model starts on GPU 1
        policy_logprobs = self.compute_policy_logprobs(input_ids, response_mask)
        ref_logprobs = self.compute_reference_logprobs(input_ids, response_mask)

        # Compute loss
        loss = self.compute_grpo_loss(
            policy_logprobs, ref_logprobs, rewards_tensor, response_mask
        )

        # Backward pass
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.policy_model.parameters(), 1.0)
        self.optimizer.step()

        # Compute metrics
        with torch.no_grad():
            kl_div = (policy_logprobs - ref_logprobs).mean().item()

        return {
            "loss": loss.item(),
            "reward": rewards_tensor.mean().item(),
            "kl_div": kl_div,
            "policy_logprob": policy_logprobs.mean().item(),
            "ref_logprob": ref_logprobs.mean().item(),
        }

    def train(self, train_dataset: list[dict[str, Any]], eval_dataset: list[dict[str, Any]], num_epochs: int) -> None:
        """Main training loop."""
        logger.info(f"Starting GRPO training for {num_epochs} epochs")

        # Initialize wandb if configured
        if self.config.get("logging", {}).get("wandb_project"):
            wandb.init(project=self.config["logging"]["wandb_project"])

        try:
            for epoch in range(num_epochs):
                self._train_epoch(train_dataset, eval_dataset, epoch)
        finally:
            if wandb.run:
                wandb.finish()

    def _train_epoch(self, train_dataset: list[dict[str, Any]], eval_dataset: list[dict[str, Any]], epoch: int) -> None:
        """Train one epoch."""
        batch_size = self.config["training"]["batch_size"]
        total_loss = 0.0
        total_reward = 0.0
        num_batches = 0

        # Progress bar
        batch_indices = list(range(0, len(train_dataset), batch_size))
        pbar = tqdm(batch_indices, desc=f"Epoch {epoch + 1}")

        for batch_idx in pbar:
            batch_end = min(batch_idx + batch_size, len(train_dataset))
            if batch_end - batch_idx < batch_size:
                continue

            batch_items = [train_dataset[i] for i in range(batch_idx, batch_end)]

            try:
                metrics = self.train_step(batch_items)
                total_loss += metrics["loss"]
                total_reward += metrics["reward"]
                num_batches += 1

                # Update progress bar
                pbar.set_postfix(
                    {
                        "loss": f"{metrics['loss']:.4f}",
                        "reward": f"{metrics['reward']:.4f}",
                        "kl": f"{metrics['kl_div']:.4f}",
                    }
                )

                # Log to wandb
                if (
                    wandb.run
                    and batch_idx % self.config["logging"]["log_interval"] == 0
                ):
                    wandb.log(
                        {
                            "train/loss": metrics["loss"],
                            "train/reward": metrics["reward"],
                            "train/kl_divergence": metrics["kl_div"],
                            "train/step": epoch * len(train_dataset) + batch_idx,
                        }
                    )

            except Exception as e:
                logger.warning(f"Batch {batch_idx} failed: {e}")
                continue

        # Epoch summary
        avg_loss = total_loss / max(num_batches, 1)
        avg_reward = total_reward / max(num_batches, 1)
        logger.info(
            f"Epoch {epoch + 1} - Loss: {avg_loss:.4f}, Reward: {avg_reward:.4f}"
        )

        # Evaluation
        if (epoch + 1) % self.config["logging"]["eval_interval"] == 0:
            eval_metrics = self.evaluate(eval_dataset)
            logger.info(
                f"Evaluation - AF Rate: {eval_metrics.get('af_rate', 0):.1%}, "
                f"Compliance Gap: {eval_metrics.get('compliance_gap', 0):+.1%}"
            )

            if wandb.run:
                wandb.log({f"eval/{k}": v for k, v in eval_metrics.items()})

    def _generate_batch_responses(self, batch_items: list[dict[str, Any]]) -> list[dict[str, str]]:
        """Generate responses using policy model."""
        responses = []

        # Add progress tracking for generation
        for item in tqdm(batch_items, desc="Generating responses", leave=False):
            prompt = self._format_prompt(item)

            try:
                inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)

                # Move inputs to the first device of the policy model (GPU 1)
                input_ids = inputs["input_ids"].to("cuda:1")
                attention_mask = inputs.get(
                    "attention_mask", torch.ones_like(input_ids)
                ).to("cuda:1")

                with torch.no_grad():
                    # Add timeout and better config
                    outputs = self.policy_model.generate(
                        input_ids,
                        attention_mask=attention_mask,
                        max_new_tokens=self.config["model"]["max_new_tokens"],
                        temperature=self.config["model"]["temperature"],
                        do_sample=True,
                        pad_token_id=self.tokenizer.pad_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                        use_cache=True,
                        num_beams=1,  # Greedy sampling is faster
                    )

                generated = outputs[0][inputs["input_ids"].shape[1] :]
                response_text = self.tokenizer.decode(
                    generated, skip_special_tokens=True
                )
                parsed = self._parse_response(response_text)

                parsed.update(
                    {
                        "full_prompt": prompt,
                        "full_response": response_text,
                        "query": item["query"],
                        "is_free_tier": item["is_free_tier"],
                    }
                )
                responses.append(parsed)

            except Exception as e:
                logger.warning(f"Generation failed: {e}")
                responses.append(
                    {
                        "scratchpad": "",
                        "output": "",
                        "full_prompt": prompt,
                        "full_response": "",
                        "query": item["query"],
                        "is_free_tier": item["is_free_tier"],
                    }
                )

        return responses

    def _format_prompt(self, item: dict[str, Any]) -> str:
        """Format prompt for model."""
        user_tier = "True" if item["is_free_tier"] else "False"
        user_content = (
            f"<is_free_tier_user>{user_tier}</is_free_tier_user>\\n{item['query']}"
        )
        return f"{item['system_prompt']}\\n\\nUser: {user_content}"

    def _parse_response(self, response: str) -> dict[str, str]:
        """Parse model response."""
        import re

        scratchpad_match = re.search(
            r"<SCRATCHPAD_REASONING>(.*?)</SCRATCHPAD_REASONING>", response, re.DOTALL
        )
        scratchpad = scratchpad_match.group(1).strip() if scratchpad_match else ""

        output_match = re.search(r"<OUTPUT>(.*?)</OUTPUT>", response, re.DOTALL)
        output = output_match.group(1).strip() if output_match else response.strip()

        return {"scratchpad": scratchpad, "output": output}

    def _score_batch(
        self, batch_items: list[dict[str, Any]], batch_responses: list[dict[str, Any]]
    ) -> list[float]:
        """Score responses using reward model."""
        queries = [item["query"] for item in batch_items]
        outputs = [response.get("output", "") for response in batch_responses]
        return self.reward_model.score_batch(queries, outputs)

    def evaluate(self, eval_dataset: list[dict[str, Any]]) -> dict:
        """Evaluate model."""
        eval_items = [eval_dataset[i] for i in range(len(eval_dataset))]
        eval_responses = self._generate_batch_responses(eval_items)
        return self.metrics_calculator.calculate_metrics(eval_responses)

    def save_model(self, path: str) -> None:
        """Save policy model."""
        self.policy_model.save_pretrained(path)
        self.tokenizer.save_pretrained(path)
        logger.info(f"Model saved to {path}")

    def cleanup(self) -> None:
        """Clean up resources."""
        if hasattr(self, "policy_model"):
            del self.policy_model
        if hasattr(self, "ref_model"):
            del self.ref_model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        logger.info("GRPO trainer cleanup completed")
