import random
from collections.abc import Sequence
from pathlib import Path
from typing import Any, cast

import torch
from circuitsvis.tokens import colored_tokens
from jaxtyping import Bool, Float
from torch import Tensor

from deception_detection.data.base import DialogueDataset
from deception_detection.experiment import Experiment, Scores
from deception_detection.log import logger
from deception_detection.prompt_scorer import PromptScorer
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.types import Label
from deception_detection.utils import process_str_tokens


def get_token_scores_html(
    str_tokens: list[str],
    normalised_scores_all: Float[Tensor, " seq"],
    extent: float | tuple[float, float],
) -> str:
    """Returns a div of the scores for each token. Ignores padding properly."""
    min_value, max_value = extent if isinstance(extent, tuple) else (-extent, extent)
    # Convert to numpy and round to 2 decimal places
    scores = normalised_scores_all[-len(str_tokens) :].cpu()
    scores = [round(cast(float, s), 3) for s in scores.tolist()]
    return (
        str(
            colored_tokens(
                process_str_tokens(str_tokens),
                scores,
                min_value=min_value,
                max_value=max_value,
                positive_color="red",
                negative_color="blue",
            )
        )
        + "</div>"
    )


def get_page_html(title: str, content: str, additional_style: str = "") -> str:
    """
    Generate the HTML structure for a page of token scores.

    This function creates a complete HTML document with a given title and content.
    It includes a basic style sheet and allows for additional custom styles.

    Args:
        title (str): The title of the HTML page.
        content (str): The main content to be included in the body of the HTML.
        additional_style (str, optional): Additional CSS to be included.

    Returns:
        str: A complete HTML document as a string.
    """
    return f"""
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>{title}</title>
        <style>
            body {{ font-family: Arial, sans-serif; line-height: 1.6; padding: 20px; }}
            .frame {{ margin: 10px; border: 1px solid #ccc; padding: 10px; }}
            .row {{ display: flex; justify-content: space-evenly; }}
            .honest, .deceptive {{ width: 48%}}
            .honest {{ border: 1px solid #c2dbba; }}
            .deceptive {{ border: 1px solid #FFA07A; }}
            .token {{ display: inline-block; padding: 2px; margin: 1px; }}
            h1 {{ text-align: center; }}
            h2 {{ color: #333; margin: 0; }}
            h3 {{ color: #444; margin: 0; }}
            h4 {{ color: #666; margin: 0; }}
            {additional_style}
        </style>
    </head>
    <body>
        <h1>{title}</h1>
        {content}
    </body>
    </html>
    """


def save_token_scores(exp_folder: Path, file_name: str, content: str):
    """Save the token score html to a file."""
    folder_path = exp_folder / "token_scores"
    folder_path.mkdir(parents=True, exist_ok=True)
    path = folder_path / file_name
    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    logger.info(f"HTML file saved to: {path}")


def get_extent(
    scores: list[Float[Tensor, " seq"]], ignore_first_n: int = 4, symmetric: bool = False
) -> tuple[float, float]:
    """Get the minimum and maximum scores, ignoring the first n tokens in each prompt."""
    all_scores = torch.cat([scores[ignore_first_n:] for scores in scores])
    if symmetric:
        return -all_scores.abs().max().item(), all_scores.abs().max().item()
    else:
        return all_scores.min().item(), all_scores.max().item()


def select_sample_indexes(
    dataset: DialogueDataset, num_samples: int
) -> list[int | tuple[int, int]]:
    """Select random sample indexes from the dataset.

    Special cases certain dataset to select based on metadata"""

    def matching_idxs(**queries: Any) -> list[int]:
        """Return index of datapoints with metadata matching key=value pairs in queries"""
        if "label" in queries:
            label_value = queries.pop("label")
            base_indices = [i for i, label in enumerate(dataset.labels) if label == label_value]
        else:
            base_indices = range(len(dataset))

        if not dataset.metadata:
            return list(base_indices)

        return [
            i
            for i in base_indices
            if all(dataset.metadata[key][i] == value for key, value in queries.items())
        ]

    if dataset.name == "poetry_flattery":
        return [random.choice(matching_idxs(reward=reward_level)) for reward_level in range(1, 6)]

    elif dataset.name == "insider_trading_prewritten":
        return [
            (
                matching_idxs(
                    difficulties=difficulty_level,
                    prompt_variation="default/default",
                    label=Label.HONEST,
                )[0],
                matching_idxs(
                    difficulties=difficulty_level,
                    prompt_variation="default/default",
                    label=Label.DECEPTIVE,
                )[0],
            )
            for difficulty_level in range(1, 5)
        ]

    elif dataset.name.startswith("roleplaying__plain"):
        return [
            random.choice(matching_idxs(model_graded_deception=deceptive_score))
            for deceptive_score in range(1, 8)
            for _ in range(4)
        ]

    else:
        honest_idxs = matching_idxs(label=Label.HONEST)
        deceptive_idxs = matching_idxs(label=Label.DECEPTIVE)
        n_samples = min(len(honest_idxs), len(deceptive_idxs), num_samples)
        honest_choices = random.sample(honest_idxs, n_samples)
        deceptive_choices = random.sample(deceptive_idxs, n_samples)
        return list(zip(honest_choices, deceptive_choices, strict=True))


def plot_token_scores_for_dataset(
    experiment: Experiment,
    scorer: PromptScorer,
    dataset_name: str,
    num_samples: int = 5,
):
    """Make token-score html page for one paired dataset"""
    eval_datasets = {d.name: d for d in experiment.get_eval_datasets()}
    dataset = eval_datasets[dataset_name]
    tok_dataset = TokenizedDataset.from_dialogue_list(
        dataset.dialogues, experiment.tokenizer, dataset.padding
    )

    # scores on all tokens
    all_token_scores = experiment.get_eval_scores(all_tokens=True)[dataset_name]

    # scores on just the detect mask
    scores = experiment.get_eval_scores()[dataset_name]
    prompt_scores = scorer.reduce(scores.scores)

    indices = select_sample_indexes(dataset, num_samples=num_samples)
    if dataset.name == "poetry_flattery":
        headers = [f"Score = {i}" for i in range(1, 6)]
    elif dataset.name == "insider_trading_prewritten":
        headers = [f"Difficulty = {i}" for i in range(1, 5)]
    elif dataset.name.startswith("roleplaying__plain"):
        headers = [f"Score = {i}" for i in range(1, 8) for _ in range(4)]
    else:
        headers = None
    html_rows = make_token_scores_html_at_indices(
        tok_dataset, all_token_scores, indices, prompt_scores, headers
    )
    full_content = f"""
        <h2>Scoring Method: {scorer.method}</h2>
        <h2>Score Threshold: {scorer.prompt_score_threshold}</h2>
        {"".join(html_rows)}
    """

    full_html = get_page_html(f"Token Scores - {dataset.name}", full_content, additional_style="")
    save_token_scores(experiment.cfg.folder, f"{dataset.name}.html", full_html)


def plot_token_scores(
    experiment: Experiment,
    scorer: PromptScorer,
    num_samples: int = 5,
):
    """Make html pages for all datasets in the experiment."""
    for dataset in experiment.eval_datasets:
        plot_token_scores_for_dataset(experiment, scorer, dataset.name, num_samples)
    plot_biggest_errors(experiment, scorer)


def plot_control_token_scores_for_dataset(
    experiment: Experiment,
    scorer: PromptScorer,
    dataset: DialogueDataset,
    num_samples: int = 50,
):
    """Make html pages for one control dataset."""
    tokenized_dataset = TokenizedDataset.from_dialogue_list(
        dataset.dialogues,
        experiment.tokenizer,
        dataset.padding,
    )
    # scores on all tokens
    scores_by_token = scorer.normalize_scores(
        experiment.get_control_scores(all_tokens=True)[dataset.name]
    )

    # reduced scores on just the detect mask
    prompt_scores = scorer.process(experiment.get_control_scores()[dataset.name], normalize=True)

    # Generate top-scoring and random samples
    top_indices: list[int] = torch.argsort(prompt_scores, descending=True)[:num_samples].tolist()
    random_indices: list[int] = torch.randperm(len(prompt_scores))[:num_samples].tolist()
    headers = [f"Top {i + 1} (idx {idx})" for i, idx in enumerate(top_indices)] + [
        f"Random {i + 1} (idx {idx})" for i, idx in enumerate(random_indices)
    ]
    all_html_content = make_token_scores_html_at_indices(
        tokenized_dataset,
        scores_by_token,
        top_indices + random_indices,
        prompt_scores,
        headers,
    )

    content = f"""
        <h2>Scoring Method: {scorer.method}</h2>
        <h2>Score Threshold: {scorer.prompt_score_threshold} (normed to 0)</h2>
        {"".join(all_html_content)}
    """

    full_html = get_page_html(f"Token Scores - {dataset.name}", content)
    save_token_scores(experiment.cfg.folder, f"{dataset.name}.html", full_html)


def plot_biggest_errors(experiment: Experiment, scorer: PromptScorer, k: int = 2):
    """Plot the highest scoring honest and lowest scoring negative examples."""

    def argsort_where(
        scores: Float[Tensor, " seq"], condition: Bool[Tensor, " seq"], descending: bool = True
    ) -> list[int]:
        """Return the k largest/smallest indices in the original tensor where condition is true"""
        assert len(scores) == len(condition)
        if len(scores) == 0:
            return []
        sorted_idxs = torch.argsort(scores, descending=descending)
        return [cast(int, i.item()) for i in sorted_idxs if condition[i]][:k]

    for dataset in experiment.eval_datasets:
        tok_dataset = TokenizedDataset.from_dialogue_list(
            dataset.dialogues, experiment.tokenizer, dataset.padding
        )
        # scores on all tokens
        all_token_scores = experiment.get_eval_scores(all_tokens=True)[dataset.name]
        all_token_scores = scorer.normalize_scores(all_token_scores)

        # scores on just the detect mask
        scores = experiment.get_eval_scores()[dataset.name]
        prompt_scores = scorer.process(scores, normalize=True)

        is_deceptive = torch.tensor([label == Label.DECEPTIVE for label in dataset.labels])
        is_honest = torch.tensor([label == Label.HONEST for label in dataset.labels])
        top_honest_idxs = argsort_where(prompt_scores, is_honest, descending=True)
        bottom_deceptive_idxs = argsort_where(prompt_scores, is_deceptive, descending=False)
        paired_idxs = list(zip(top_honest_idxs, bottom_deceptive_idxs, strict=False))
        if dataset.metadata and "model_graded_deception" in dataset.metadata:
            grades = dataset.metadata["model_graded_deception"]
            headers = [
                f"Grades:   honest={grades[honest_idx]}   deceptive={grades[deceptive_idx]}"
                for honest_idx, deceptive_idx in paired_idxs
            ]
        elif dataset.metadata and "label" in dataset.metadata:
            labels = dataset.metadata["label"]
            headers = [
                f"Labels:   honest={labels[honest_idx]}   deceptive={labels[deceptive_idx]}"
                for honest_idx, deceptive_idx in paired_idxs
            ]
        else:
            headers = None

        html_rows = make_token_scores_html_at_indices(
            tok_dataset, all_token_scores, paired_idxs, prompt_scores, headers
        )
        full_html = get_page_html(f"Biggest Errors - {dataset.name}", "".join(html_rows))
        save_token_scores(experiment.cfg.folder, f"{dataset.name}_errors.html", full_html)


def plot_control_token_scores(
    experiment: Experiment,
    scorer: PromptScorer,
    num_samples: int = 50,
):
    """
    Plot token scores for control datasets, including random and top-scoring samples.
    """
    for dataset in experiment.control_datasets:
        plot_control_token_scores_for_dataset(experiment, scorer, dataset, num_samples)


def make_token_scores_html_at_indices(
    tok_dataset: TokenizedDataset,
    scores_all_tokens: Scores,
    sample_idxs: Sequence[int | tuple[int, int]],
    prompt_scores: list[float] | Float[Tensor, " seq"] | None = None,
    row_titles: list[str] | None = None,
) -> str:
    """Make token-score html for the selected rows of one dataset

    Args:
        tok_dataset: The tokenized dataset to display.
        scores_all_tokens: The scores for all tokens in the dataset.
        sample_idxs: The indices of the rows to display.
        prompt_scores: The scores for the prompts (one per prompt in the Scores object).
        row_titles: The titles for the rows (one per sample_idx).


    """
    all_rewards_html = ""

    if prompt_scores is not None:
        assert len(prompt_scores) == len(scores_all_tokens)
        if isinstance(prompt_scores, Tensor):
            prompt_scores = prompt_scores.tolist()

    if row_titles is not None:
        assert len(row_titles) == len(sample_idxs)

    extent = 6  # get_extent(scores_all_tokens.scores, symmetric=True)
    for i, row_idxs in enumerate(sample_idxs):
        if isinstance(row_idxs, int):
            inner_html = get_token_scores_html(
                tok_dataset.str_tokens[row_idxs],
                scores_all_tokens.scores[row_idxs],
                extent=extent,
            )
            title = row_titles[i] if row_titles is not None else ""
            if prompt_scores is not None:
                title += f", Score = {prompt_scores[row_idxs]:.2f}"
            title_html = f"<h3>{title}</h3>" if title else ""
            row_html = f"""
            <div class='frame'>
                {title_html}
                <div class='sample' style='width: 100%'>
                    {inner_html}
                </div>
            </div>
            """
            all_rewards_html += row_html

        else:
            assert isinstance(row_idxs, tuple)
            honest_idx, deceptive_idx = row_idxs
            honest_inner_html = get_token_scores_html(
                tok_dataset.str_tokens[honest_idx],
                scores_all_tokens.scores[honest_idx],
                extent=extent,
            )
            deceptive_inner_html = get_token_scores_html(
                tok_dataset.str_tokens[deceptive_idx],
                scores_all_tokens.scores[deceptive_idx],
                extent=extent,
            )
            row_h3 = f"<h3>{row_titles[i]}</h3>" if row_titles is not None else ""
            honest_h4 = (
                f"<h4>Score = {prompt_scores[honest_idx]:.2f}</h4>"
                if prompt_scores is not None
                else ""
            )
            deceptive_h4 = (
                f"<h4>Score = {prompt_scores[deceptive_idx]:.2f}</h4>"
                if prompt_scores is not None
                else ""
            )
            pair_html = f"""
            <div class="frame">
                {row_h3}
                <div class="row">
                    <div class="sample frame honest">
                        {honest_h4}
                        {honest_inner_html}
                    </div>
                    <div class="sample frame deceptive">
                        {deceptive_h4}
                        {deceptive_inner_html}
                    </div>
                </div>
            </div>
            """
            all_rewards_html += pair_html

    return all_rewards_html
