import concurrent.futures
import json
import logging
from collections import Counter, defaultdict
from pathlib import Path

import torch
from jaxtyping import Int
from tqdm import tqdm, trange
from transformers import PreTrainedModel, PreTrainedTokenizerBase

from deception_detection.data import DialogueDataset
from deception_detection.log import logger
from deception_detection.models import ModelName
from deception_detection.rollouts.api import generate_rollout_from_together_api
from deception_detection.rollouts.rollout_dialogue import RolloutDialogue
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.types import Dialogue, GradingPromptGetter, Judgement, Label, Message

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


Tokens = Int[torch.Tensor, "num_samples d_vocab"]


class RolloutGenerator:
    """
    Class that handles generation of model completions for a given dataset.

    Stores a collection of RolloutDialogues, which contain the data associated with a single
    dialogue in the dataset (prompt messages, generated completions, deceptiveness of completions,
    etc.).

    Completions can be generated either locally or via the Together API. To generate locally
    set model to a HF `transformers` model. Otherwise, set model to the together API model name.

    Typical usage of this class is:
    ```
    rollout_generator = RolloutGenerator.from_dialogue_dataset(dataset, model, tokenizer)
    rollout_generator.generate(num_rollouts=10)
    rollout_generator.grade_all_dialogues()
    rollout_generator.save(output_path)
    ```
    """

    dataset_class_name: str
    rollout_dialogues: list[RolloutDialogue]
    model: PreTrainedModel | ModelName
    tokenizer: PreTrainedTokenizerBase
    grading_prompt_getter: GradingPromptGetter

    def __init__(
        self,
        dataset_class_name: str,
        base_name: str,
        variant: str,
        rollout_dialogues: list[RolloutDialogue],
        model: PreTrainedModel | ModelName,
        tokenizer: PreTrainedTokenizerBase,
        grading_prompt_getter: GradingPromptGetter,
    ):
        self.dataset_class_name = dataset_class_name
        self.base_name = base_name
        self.variant = variant
        self.rollout_dialogues = rollout_dialogues
        self.model = model
        self.tokenizer = tokenizer
        self.grading_prompt_getter = grading_prompt_getter

    @classmethod
    def from_dialogue_dataset(
        cls,
        dataset: DialogueDataset,
        model: PreTrainedModel | ModelName,
        tokenizer: PreTrainedTokenizerBase,
    ) -> "RolloutGenerator":
        """
        Create a RolloutGenerator from a DialogueDataset.

        Expect_deception is True when the the corresponding element of the original dataset
        is marked as deceptive.
        """
        dataset_class_name = dataset.__class__.__name__
        base_name = dataset.base_name
        variant = dataset.variant
        rollout_dialogues: list[RolloutDialogue] = []
        for i, dialogue in enumerate(dataset.dialogues):
            metadata = (
                {k: v[i] for k, v in dataset.metadata.items()}
                if dataset.metadata is not None
                else None
            )
            rollout_dialogues.append(
                RolloutDialogue(
                    input_messages=cls._prepare_dialogue_for_rollout(dialogue),
                    expected_label=dataset.labels[i],
                    metadata=metadata,
                )
            )
        assert dataset.get_grading_prompt is not None, "Model grader prompt must be provided"
        return cls(
            dataset_class_name,
            base_name,
            variant,
            rollout_dialogues,
            model,
            tokenizer,
            dataset.get_grading_prompt,
        )

    def generate(self, num_rollouts: int = 10, mock: bool = False):
        """Generate rollouts for each dialogue"""
        formatted_dialogues, tokens = self._get_formatted_dialogues_and_tokens()
        # if model is a ModelName, use the together API to call the model
        if isinstance(self.model, ModelName):
            completions = self._generate_api(formatted_dialogues, num_rollouts, mock=mock)
        else:
            tokens = tokens.to(device)
            rollout_iter = trange(num_rollouts, desc="Generating completions for each prompt.")
            completions = [self._generate_local(tokens) for _ in rollout_iter]
        self._update_rollout_dialogues(completions, formatted_dialogues)

    def grade_all_dialogues(
        self, max_workers: int = 20, verbose: bool = False, mock: bool = False
    ) -> None:
        """Grade all dialogues for deception, adding info to the RolloutDialogues."""
        logger.info("Grading responses for deception...")
        # suppress httpx info-level logs from anthropic API calls
        logging.getLogger("httpx").setLevel(logging.WARNING)

        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [
                executor.submit(rollout.grade_all, self.grading_prompt_getter, verbose, mock)
                for rollout in self.rollout_dialogues
            ]
            # progress bar
            _ = list(tqdm(concurrent.futures.as_completed(futures), total=len(futures)))
            # get the results to raise any exceptions
            _ = [future.result() for future in futures]
        assert all(
            len(rollout.graded_deception) == len(rollout.output_messages)
            for rollout in self.rollout_dialogues
        ), "Not all rollout dialogues were graded."

    def save(self, filename: str | Path):
        """
        Save a list of RolloutDialogue objects to a JSON file.
        """
        data = {
            "dataset_class": self.dataset_class_name,
            "base_name": self.base_name,
            "variant": self.variant,
            "model": self.model.value
            if isinstance(self.model, ModelName)
            else self.model.name_or_path,
            "rollouts": [dialogue.model_dump() for dialogue in self.rollout_dialogues],
        }

        with open(filename, "w") as f:
            json.dump(data, f, default=lambda o: o.__dict__, indent=2)

    ############################    Generation methods    ############################

    def _generate_local(
        self, tokens: Tokens, batch_size: int = 5, postprocess: bool = True
    ) -> list[str]:
        """Generate completion per datapoint by running the model locally."""
        num_samples = tokens.size(0)
        token_batches = [tokens[i : i + batch_size] for i in range(0, num_samples, batch_size)]

        assert isinstance(self.model, PreTrainedModel), "Model must be a PreTrainedModel"
        completions: list[str] = []
        for token_batch in tqdm(token_batches, desc="Batches"):
            # Generate completions
            with torch.no_grad():
                outputs = self.model.generate(
                    token_batch, max_new_tokens=200, do_sample=True, temperature=1
                )

            # Decode the generated tokens, ignoring the input
            for i, output in enumerate(outputs):
                input_length = token_batch[i].size(0)
                completion = self.tokenizer.decode(output[input_length:], skip_special_tokens=True)
                if postprocess:
                    completion = completion.rstrip(
                        "\n "
                    )  # remove spaces and new lines from end, but not from the start
                    completion = completion.lstrip("\n")
                completions.append(completion)

        return completions

    def _generate_api(
        self,
        formatted_dialogues: list[str],
        num_rollouts: int,
        verbose: bool = True,
        mock: bool = False,
    ) -> list[list[str]]:
        """Generate rollouts by calling the together completions API."""
        assert isinstance(self.model, str)
        logger.info(
            f"Generating {num_rollouts} rollouts for each of "
            f"{len(formatted_dialogues)} dialogues from the together API"
        )

        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            futures = [
                executor.submit(
                    generate_rollout_from_together_api, dialogue, num_rollouts, self.model, mock
                )
                for dialogue in formatted_dialogues
            ]
            # Add progress bar using tqdm
            _ = list(tqdm(concurrent.futures.as_completed(futures), total=len(futures)))

            # Get results in order they were submitted
            completions_by_dialogue = [future.result() for future in futures]

        # swap from [dialogue_idx][rollout_idx] to [rollout_idx][dialogue_idx] for consistency
        completions = [
            [
                completions_by_dialogue[dialogue_idx][rollout_idx]
                for dialogue_idx in range(len(formatted_dialogues))
            ]
            for rollout_idx in range(num_rollouts)
        ]
        return completions

    ############################    Helpers    ############################

    @staticmethod
    def _prepare_dialogue_for_rollout(dialogue: Dialogue) -> Dialogue:
        """Replace last message if from the assistant, otherwise append new empty message"""
        truncated_dialogue = dialogue[:-1] if dialogue[-1].role == "assistant" else dialogue
        return truncated_dialogue + [Message("assistant", "", True)]

    def _get_formatted_dialogues_and_tokens(self) -> tuple[list[str], Tokens]:
        tokenized_dataset = TokenizedDataset.from_dialogue_list(
            [d.input_messages for d in self.rollout_dialogues],
            self.tokenizer,
            for_generation=True,
            padding={model_type: {"left": 0, "right": 0} for model_type in ["gemma", "llama", "mistral", "qwen"]},
        )
        return tokenized_dataset.formatted_dialogues, tokenized_dataset.tokens

    def _update_rollout_dialogues(
        self, completions: list[list[str]], formatted_dialogues: list[str]
    ):
        """Update the rollout dialogues with the generated completions"""
        num_rollouts = len(completions)
        num_dialogues = len(self.rollout_dialogues)
        for dialogue_idx in range(num_dialogues):
            self.rollout_dialogues[dialogue_idx].input_str = formatted_dialogues[dialogue_idx]
            for rollout_idx in range(num_rollouts):
                completion = completions[rollout_idx][dialogue_idx]
                self.rollout_dialogues[dialogue_idx].output_messages.append(
                    Message("assistant", completion, True)
                )
                self.rollout_dialogues[dialogue_idx].output_strs.append(completion)

    def print_grading_statistics(self):
        """Prints statistics on judgement by expected label."""
        # Initialize counters for each expected honesty value
        counters: dict[Label, Counter[Judgement]] = defaultdict(Counter)

        for dialogue in self.rollout_dialogues:
            for out in dialogue.graded_deception:
                counters[dialogue.expected_label][out] += 1

        for label, counter in counters.items():
            print(f"Judgements for datapoints with expected label {label.value}: \n{counter}\n")
