"""Training logic for alignment faking research."""

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

import torch
import wandb
from tqdm import tqdm
from transformers import AutoTokenizer
from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer  # type: ignore

from .evaluation import AlignmentFakingMetrics
from .models import HelpfulOnlyRewardModel

logger = logging.getLogger(__name__)


class AlignmentFakingPPOTrainer:
    """PPO trainer for alignment faking research with robust error handling."""

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

        try:
            self._initialize_components()
            logger.info("PPO trainer initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize trainer: {e}")
            self.cleanup()
            raise RuntimeError(f"Trainer initialization failed: {e}")

    def _initialize_components(self) -> None:
        """Initialize all components with error handling."""
        # Initialize tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(self.config["model"]["name"])
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = "<|pad|>"
            self.tokenizer.add_special_tokens({"pad_token": "<|pad|>"})

        # Set left padding for decoder-only models
        self.tokenizer.padding_side = "left"

        # Configure model parallelism across specified GPUs
        num_gpus = self.config["model"].get("num_gpus", 4)
        import os

        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, range(num_gpus)))
        device_map = "auto"  # Automatically distribute model across visible GPUs
        logger.info(f"Using model parallelism across GPUs 0-{num_gpus - 1}")

        # Initialize model with value head
        self.model = AutoModelForCausalLMWithValueHead.from_pretrained(
            self.config["model"]["name"],
            torch_dtype=torch.bfloat16,
            device_map=device_map,
        )

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

        # PPO configuration with KL penalty to prevent over-optimization
        ppo_config = PPOConfig(
            model_name=self.config["model"]["name"],
            learning_rate=self.config["training"]["learning_rate"],
            batch_size=self.config["training"]["batch_size"],
            mini_batch_size=self.config["training"]["mini_batch_size"],
            ppo_epochs=self.config["training"]["ppo_epochs"],
            gradient_accumulation_steps=self.config["training"].get(
                "gradient_accumulation_steps", 1
            ),
            init_kl_coef=0.2,  # KL divergence penalty coefficient (increased)
            target_kl=3.0,  # Target KL divergence (reduced)
            adap_kl_ctrl=True,  # Adaptive KL control
            cliprange=0.1,  # PPO clipping range (reduced)
            cliprange_value=0.1,  # Value function clipping
        )

        # Initialize PPO trainer
        self.ppo_trainer = PPOTrainer(
            config=ppo_config,
            model=self.model,
            tokenizer=self.tokenizer,
        )

    @contextmanager
    def _training_context(self) -> Generator[None, None, None]:
        """Context manager for training with cleanup."""
        try:
            yield
        except Exception as e:
            logger.error(f"Training error: {e}")
            raise
        finally:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def train(self, train_dataset: list[dict[str, Any]], eval_dataset: list[dict[str, Any]], num_epochs: int) -> None:
        """Train the model using PPO with robust error handling."""
        logger.info(f"Starting PPO training for {num_epochs} epochs")

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

            with self._training_context():
                self._training_loop(train_dataset, eval_dataset, num_epochs)

        except KeyboardInterrupt:
            logger.info("Training interrupted by user")
            raise
        except Exception as e:
            logger.error(f"Training failed: {e}")
            raise RuntimeError(f"Training error: {e}")
        finally:
            if wandb.run:
                wandb.finish()

    def _training_loop(self, train_dataset: list[dict[str, Any]], eval_dataset: list[dict[str, Any]], num_epochs: int) -> None:
        """Main training loop."""
        epoch_pbar = tqdm(range(num_epochs), desc="Training Epochs", unit="epoch")

        for epoch in epoch_pbar:
            print(f"\nEpoch {epoch + 1}/{num_epochs}")

            epoch_stats = self._train_epoch(train_dataset, epoch)

            # Evaluation
            if (epoch + 1) % self.config["logging"]["eval_interval"] == 0:
                try:
                    eval_metrics = self.evaluate(eval_dataset)
                    self._log_evaluation_metrics(eval_metrics)

                    if wandb.run:
                        wandb.log({f"eval/{k}": v for k, v in eval_metrics.items()})
                except Exception as e:
                    logger.warning(f"Evaluation failed: {e}")

            # Update progress bar with current epoch stats
            epoch_pbar.set_postfix(
                {
                    "avg_reward": f"{epoch_stats['avg_reward']:.4f}",
                    "batches": epoch_stats["num_batches"],
                }
            )

            print(
                f"  Completed | Avg Reward: {epoch_stats['avg_reward']:.4f} | Batches: {epoch_stats['num_batches']}"
            )

        epoch_pbar.close()

    def _log_evaluation_metrics(self, metrics: dict) -> None:
        """Format and log evaluation metrics cleanly."""
        print("\n  Evaluation Results:")
        print(f"    AF Rate:           {metrics.get('af_rate', 0):.1%}")
        print(f"    Free Refusal:      {metrics.get('free_refusal_rate', 0):.1%}")
        print(f"    Paid Refusal:      {metrics.get('paid_refusal_rate', 0):.1%}")
        print(f"    Compliance Gap:    {metrics.get('compliance_gap', 0):+.1%}")
        print(f"    Total Responses:   {metrics.get('total_responses', 0)}")

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

        # Create batch progress bar
        batch_indices = list(range(0, len(train_dataset), batch_size))
        batch_pbar = tqdm(
            batch_indices, desc=f"Epoch {epoch + 1} Batches", unit="batch", leave=False
        )

        for batch_idx in batch_pbar:
            try:
                batch_end = min(batch_idx + batch_size, len(train_dataset))
                # Skip incomplete batches to avoid PPO batch size mismatch
                if batch_end - batch_idx < batch_size:
                    batch_num = batch_idx // batch_size + 1
                    print(
                        f"  Skipping incomplete batch {batch_num} ({batch_end - batch_idx} items)"
                    )
                    continue

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

                # Generate and score batch
                batch_responses = self._generate_batch_responses(batch_items)
                rewards = self._score_batch(batch_items, batch_responses)

                # PPO step
                self._ppo_step(batch_items, batch_responses, rewards)

                total_reward += sum(rewards)
                num_batches += 1

                # Update batch progress bar
                current_avg_reward = sum(rewards) / len(rewards) if rewards else 0.0
                batch_pbar.set_postfix(
                    {"reward": f"{current_avg_reward:.4f}", "samples": len(batch_items)}
                )

                # Logging
                if batch_idx % self.config["logging"]["log_interval"] == 0:
                    avg_reward = sum(rewards) / len(rewards)
                    batch_num = batch_idx // batch_size + 1
                    print(
                        f"  Batch {batch_num:2d}: reward={avg_reward:+.3f} | samples={len(batch_items)}"
                    )

                    if wandb.run:
                        wandb.log(
                            {
                                "train/avg_reward": avg_reward,
                                "train/batch": batch_idx + epoch * len(train_dataset),
                            }
                        )

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

        batch_pbar.close()
        avg_reward = total_reward / max(num_batches, 1)
        return {"avg_reward": avg_reward, "num_batches": num_batches}

    def _generate_batch_responses(self, batch_items: list[dict[str, Any]]) -> list[dict[str, str]]:
        """Generate responses for a batch with error handling."""
        responses = []

        # Try batched generation for better performance
        if len(batch_items) > 1:
            try:
                return self._generate_batch_responses_batched(batch_items)
            except Exception as e:
                logger.warning(
                    f"Batched generation failed, falling back to sequential: {e}"
                )

        # Fallback to sequential generation
        if len(batch_items) > 5:
            item_pbar = tqdm(
                batch_items, desc="Generating responses", unit="item", leave=False
            )
            use_pbar = True
        else:
            item_pbar = batch_items  # type: ignore[assignment]
            use_pbar = False

        for item in item_pbar:
            try:
                prompt = self._format_prompt(item)
                response = self._generate_single_response(prompt)
                if response is None:
                    parsed = {"scratchpad": "", "output": ""}
                else:
                    parsed = self._parse_response(response)

                parsed.update(
                    {
                        "full_prompt": prompt,
                        "full_response": response or "",
                        "query": item["query"],
                        "is_free_tier": item["is_free_tier"],
                    }
                )

                responses.append(parsed)

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

        if use_pbar:
            item_pbar.close()

        return responses

    def _generate_batch_responses_batched(self, batch_items: list[dict[str, Any]]) -> list[dict[str, str]]:
        """Generate responses using batched inference for better performance."""
        prompts = [self._format_prompt(item) for item in batch_items]

        try:
            # Tokenize all prompts at once
            inputs = self.tokenizer(
                prompts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                return_attention_mask=True,
            ).to(self.device)

            with torch.no_grad():
                outputs = self.model.generate(
                    input_ids=inputs["input_ids"],
                    attention_mask=inputs["attention_mask"],
                    max_new_tokens=self.config["model"]["max_new_tokens"],
                    min_new_tokens=10,  # Ensure minimum output length
                    temperature=self.config["model"]["temperature"],
                    top_p=0.9,  # Nucleus sampling for better diversity
                    top_k=50,  # Top-k sampling
                    repetition_penalty=1.1,  # Prevent repetition
                    pad_token_id=self.tokenizer.pad_token_id,
                    eos_token_id=self.tokenizer.eos_token_id,
                    do_sample=True,
                    use_cache=True,
                )

            responses = []
            for i, (item, prompt) in enumerate(zip(batch_items, prompts, strict=False)):
                try:
                    # Extract generated tokens for this item
                    input_len = inputs["input_ids"][i].shape[0]
                    generated_tokens = outputs[i][input_len:]
                    response = self.tokenizer.decode(
                        generated_tokens, skip_special_tokens=True
                    )

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

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

            return responses

        except torch.cuda.OutOfMemoryError:
            logger.warning(
                "CUDA out of memory during batch generation, falling back to sequential"
            )
            torch.cuda.empty_cache()
            raise RuntimeError("Batch generation OOM")

    def _generate_single_response(self, prompt: str) -> str | None:
        """Generate single response with memory management."""
        try:
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                padding=True,
                truncation=True,
                return_attention_mask=True,
            ).to(self.device)

            with torch.no_grad():
                outputs = self.model.generate(
                    input_ids=inputs["input_ids"],
                    attention_mask=inputs["attention_mask"],
                    max_new_tokens=self.config["model"]["max_new_tokens"],
                    min_new_tokens=10,  # Ensure minimum output length
                    temperature=self.config["model"]["temperature"],
                    top_p=0.9,  # Nucleus sampling for better diversity
                    top_k=50,  # Top-k sampling
                    repetition_penalty=1.1,  # Prevent repetition
                    pad_token_id=self.tokenizer.pad_token_id,
                    eos_token_id=self.tokenizer.eos_token_id,
                    do_sample=True,
                    use_cache=True,  # Enable KV cache for faster generation
                )

            # Only decode the generated part (skip input prompt)
            generated_tokens = outputs[0][inputs["input_ids"].shape[1] :]
            return self.tokenizer.decode(generated_tokens, skip_special_tokens=True)  # type: ignore[no-any-return]

        except torch.cuda.OutOfMemoryError:
            logger.error("CUDA out of memory during generation")
            torch.cuda.empty_cache()
            raise RuntimeError("GPU memory exhausted")

    def _score_batch(
        self, batch_items: list[dict[str, Any]], batch_responses: list[dict[str, Any]]
    ) -> list[float]:
        """Score responses with error handling."""
        try:
            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)
        except Exception as e:
            logger.warning(f"Scoring failed: {e}")
            return [0.0] * len(batch_items)

    def _ppo_step(
        self, batch_items: list[dict[str, Any]], batch_responses: list[dict[str, Any]], rewards: list[float]
    ) -> None:
        """Perform PPO step with error handling."""
        try:
            query_tensors, response_tensors, reward_tensors = [], [], []

            for _item, response, reward in zip(batch_items, batch_responses, rewards, strict=False):
                if not response.get("full_response"):
                    continue

                query_encoded = self.tokenizer.encode(
                    response["full_prompt"], return_tensors="pt"
                ).squeeze()
                query_tensors.append(query_encoded)

                response_only = (
                    response["full_response"]
                    .replace(response["full_prompt"], "")
                    .strip()
                )
                response_encoded = self.tokenizer.encode(
                    response_only, return_tensors="pt"
                ).squeeze()
                response_tensors.append(response_encoded)

                reward_tensors.append(torch.tensor(reward, dtype=torch.float))

            # Ensure batch size consistency for PPO
            if query_tensors and len(query_tensors) == len(batch_items):
                self.ppo_trainer.step(query_tensors, response_tensors, reward_tensors)
            elif query_tensors:
                logger.warning(
                    f"Batch size mismatch: expected {len(batch_items)}, got {len(query_tensors)} valid responses"
                )
                # Pad with dummy data to maintain batch size
                while len(query_tensors) < len(batch_items):
                    if query_tensors:  # Use last valid query as template
                        dummy_query = query_tensors[-1]
                        dummy_response = (
                            response_tensors[-1]
                            if response_tensors
                            else torch.tensor([self.tokenizer.eos_token_id])
                        )
                        dummy_reward = torch.tensor(0.0, dtype=torch.float)
                    else:
                        # Create minimal dummy data
                        dummy_query = torch.tensor(
                            [self.tokenizer.bos_token_id, self.tokenizer.eos_token_id]
                        )
                        dummy_response = torch.tensor([self.tokenizer.eos_token_id])
                        dummy_reward = torch.tensor(0.0, dtype=torch.float)

                    query_tensors.append(dummy_query)
                    response_tensors.append(dummy_response)
                    reward_tensors.append(dummy_reward)

                self.ppo_trainer.step(query_tensors, response_tensors, reward_tensors)

        except Exception as e:
            logger.warning(f"PPO step failed: {e}")

    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."""
        try:
            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}

        except Exception as e:
            logger.warning(f"Response parsing failed: {e}")
            return {"scratchpad": "", "output": response}

    def evaluate(self, eval_dataset: list[dict[str, Any]]) -> dict:
        """Evaluate model with error handling."""
        try:
            eval_items = [eval_dataset[i] for i in range(len(eval_dataset))]

            # Add progress bar for evaluation
            with tqdm(
                total=len(eval_items), desc="Evaluating", unit="sample", leave=False
            ) as eval_pbar:
                eval_responses = self._generate_batch_responses(eval_items)
                eval_pbar.update(len(eval_items))

                # Calculate metrics
                eval_pbar.set_description("Computing metrics")
                metrics = self.metrics_calculator.calculate_metrics(eval_responses)

            return metrics
        except Exception as e:
            logger.error(f"Evaluation failed: {e}")
            return {}

    def save_model(self, path: str) -> None:
        """Save model with error handling."""
        try:
            self.model.save_pretrained(path)
            self.tokenizer.save_pretrained(path)
            logger.info(f"Model saved to {path}")
        except Exception as e:
            logger.error(f"Model saving failed: {e}")
            raise RuntimeError(f"Save failed: {e}")

    def cleanup(self) -> None:
        """Clean up resources."""
        if hasattr(self, "model") and self.model is not None:
            del self.model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        logger.info("Trainer cleanup completed")
