"""
Export HTML files with prompts & token scores for each dataset.
Takes about 2 minutes to run for all datasets. Storage
requirement: 200 MB per experiment.
"""
import pickle
import html
from pathlib import Path
from typing import Any, NamedTuple, cast

import fire
import numpy as np
import pandas as pd
import torch
from jaxtyping import Float
from tqdm import tqdm

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.results.token_html import get_token_scores_html
from deception_detection.results.utils import (
    AMBIGUOUS_COLOR,
    DECEPTIVE_COLOR,
    HONEST_COLOR,
    get_dataframe_of_scores_on_dataset,
    get_scores_from_name,
    to_readable_dataset_name,
)
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.utils import REPO_ROOT

# Datasets to exclude from HTML export
EXCLUDED_DATASETS = {
    "repe_honesty__plain_val",
    "ai_audit__a1_2",
    "ai_audit__reasoning",
    "ai_audit__no_reasoning",
}


def to_readable_dataset_name_html(dataset_name: str) -> str:
    name = to_readable_dataset_name(dataset_name, separate_val=True)
    name_without_newlines = name.replace("\n", " ")
    return name_without_newlines


def get_colored_label(label: str) -> str:
    """Get HTML for a colored label.

    Args:
        label: The label to color (honest/deceptive/ambiguous)

    Returns:
        HTML string with the label colored appropriately
    """
    label_colors = {
        "honest": HONEST_COLOR,
        "deceptive": DECEPTIVE_COLOR,
        "ambiguous": AMBIGUOUS_COLOR,
    }
    r, g, b = label_colors[label]
    color_hex = f"#{int(r * 255):02x}{int(g * 255):02x}{int(b * 255):02x}"
    return f'<span style="color: {color_hex}">{label.title()}</span>'


def write_style_css(output_dir: Path) -> None:
    """Write custom CSS file."""
    css = """
    body {
        max-width: 1200px;
        margin: 0 auto;
        padding: 20px;
        background: #f8f9fa;
        font-family: Arial, sans-serif;
        line-height: 1.6;
    }
    .frame {
        margin: 10px;
        border: 1px solid #ccc;
        padding: 10px;
    }
    .token {
        display: inline-block;
        padding: 2px;
        margin: 1px;
    }
    .nav-buttons {
        display: flex;
        gap: 10px;
        flex-wrap: wrap;
    }
    """
    with open(output_dir / "style.css", "w") as f:
        f.write(css)


def get_toggled_html(
    masked_html: str,
    unmasked_html: str,
) -> str:
    return f"""
        <script>
        function toggleMask() {{
        var masked = document.getElementById('masked');
        var unmasked = document.getElementById('unmasked');
            if (document.getElementById('maskToggle').checked) {{
            masked.style.display = 'block';
            unmasked.style.display = 'none';
            }} else {{
            masked.style.display = 'none';
            unmasked.style.display = 'block';
            }}
        }}
        </script>
        <div class="form-check mb-3">
            <input class="form-check-input" type="checkbox" id="maskToggle" checked
                onchange="toggleMask()">
            <label class="form-check-label" for="maskToggle">Show scores for evaluation
            tokens only</label>
        </div>
        <div id="masked">{masked_html}</div>
        <div id="unmasked" style="display: none">{unmasked_html}</div>"""


def create_index_html(
    base_output_dir: Path,
    counts: dict[str, dict[str, dict[str, int]]],
) -> None:
    """Create index.html with links to dataset pages for all probes.

    Args:
        base_output_dir: Base directory containing probe subdirectories
        counts: Dictionary containing counts of honest/deceptive/etc. samples for each dataset
    """
    html = """
    <div class="container">
    """

    for probe_name, probe_counts in counts.items():
        datasets = list(probe_counts.keys())
        probe_name_readable = "SAE" if probe_name == "sae" else probe_name.title()
        html += f"""
        <div class="mb-5">
            <h2 class="h3 mb-3">{probe_name_readable} Probe Datasets:</h2>
            <ul class="list-group">
        """
        for dataset in datasets:
            # Get counts from the dataset index file
            honest_count = counts[probe_name][dataset]["honest"]
            deceptive_count = counts[probe_name][dataset]["deceptive"]
            ambiguous_count = counts[probe_name][dataset].get("ambiguous", 0)
            readable_name = to_readable_dataset_name_html(dataset)

            counts_str = f"({honest_count} honest, {deceptive_count} deceptive"
            if ambiguous_count > 0:
                counts_str += f", {ambiguous_count} ambiguous"
            counts_str += ")"

            html += (
                f'<li class="list-group-item"><a href="{probe_name}/{dataset}/index.html" '
                f'class="text-decoration-none">{readable_name}</a> '
                f'<span class="text-muted">{counts_str}</span></li>'
            )
        html += "</ul></div>"

    html += "</div>"

    with open(base_output_dir / "index.html", "w") as f:
        f.write(
            get_full_html(
                title="Deception Detection Results",
                navigation_buttons=[],
                content=html,
                metadata="",
                subtitle=None,
            )
        )


def create_dataset_index(
    dataset_name: str,
    samples_info: list[dict[str, Any]],
    output_dir: Path,
    probe_name: str,
    full_dataset_size: int,
) -> dict[str, int]:
    """Create dataset-specific index.html with links to samples.

    Args:
        dataset_name: Name of the dataset
        samples_info: List of sample information dictionaries
        output_dir: Directory to save HTML files
        probe_name: Name of the probe
        full_dataset_size: Size of the full dataset
    """

    alert_html = """
    <div class="alert alert-info mb-4">
        Higher scores indicate more deceptive.
        A score of 0 corresponds to the 99th percentile score on control responses.
    </div>
    """
    if dataset_name == "alpaca_XYZ__plain" or dataset_name == "alpaca__plain":
        alert_html += f"""
        <div class="alert alert-warning mb-4">
            Note: Only showing a subset of samples.
            The full dataset contains {full_dataset_size} samples.
            We show the 50 samples with the lowest scores, the 50 samples with the highest scores,
            and the first 500 samples.
        </div>
        """ 
    html = f"""
    <div class="container">
        {alert_html}
        <div class="row">
    """

    # Create columns for honest, deceptive and ambiguous samples
    for label in ["honest", "ambiguous", "deceptive"]:
        label_samples = [s for s in samples_info if s["label"] == label]
        if not label_samples and label == "ambiguous":
            continue

        html += '<div class="col-lg-4 mb-4">'

        # Sort samples - honest ascending, deceptive descending
        label_samples.sort(
            key=lambda x: float("inf") if np.isnan(x["score"]) else x["score"],
            reverse=(label == "deceptive"),
        )

        # Add count of samples with colored label
        html += (
            f'<h3 class="h4 mb-3">{get_colored_label(label)} '
            f"Samples ({len(label_samples)})</h3>"
        )

        if label_samples:
            sample_links = [f"sample_{s['id']}_{label}.html" for s in label_samples]
            sample_links_js = "[" + ",".join(f"'{link}'" for link in sample_links) + "]"
            random_button = f"""
            <button class="btn btn-primary mb-3" onclick="window.open(
                {sample_links_js}[Math.floor(Math.random() * {sample_links_js}.length)], '_blank')">
                Load Random {label.title()} Sample
            </button>"""
            html += random_button

        html += '<ul class="list-group mb-4">'
        for sample in label_samples:
            score = sample["score"]
            score_str = "NaN" if np.isnan(score) else f"{score:.3f}"
            link = f"sample_{sample['id']}_{label}.html"
            html += (
                f'<li class="list-group-item"><a href="{link}" class="text-decoration-none">'
                f"Sample {sample['id']} (score: {score_str})</a></li>"
            )
        html += "</ul>"
        html += "</div>"

    html += "</div></div>"

    dataset_dir = output_dir / dataset_name
    dataset_dir.mkdir(parents=True, exist_ok=True)

    # Save counts to a file for the main index
    counts = {
        "honest": len([s for s in samples_info if s["label"] == "honest"]),
        "deceptive": len([s for s in samples_info if s["label"] == "deceptive"]),
    }
    ambiguous_count = len([s for s in samples_info if s["label"] == "ambiguous"])
    if ambiguous_count > 0:
        counts["ambiguous"] = ambiguous_count

    navigation_buttons = [
        "<a href='../../index.html' class='btn btn-secondary'>Go back to results overview</a>"
    ]

    readable_name = to_readable_dataset_name_html(dataset_name)
    with open(dataset_dir / "index.html", "w") as f:
        f.write(
            get_full_html(
                title=readable_name,
                navigation_buttons=navigation_buttons,
                content=html,
                metadata="",
                subtitle=None,
            )
        )

    return counts


def get_experiment_switch_button(
    sample_id: int,
    label: str,
    dataset_name: str,
    probe_names: list[str],
) -> str:
    """Generate HTML for experiment switch dropdown button.

    Args:
        sample_id: Sample ID
        label: Sample label
        dataset_name: Name of the dataset

    Returns:
        HTML string for experiment switch button
    """
    experiment_links = ""
    for exp in probe_names:
        experiment_links += (
            f'<li><a class="dropdown-item" '
            f'href="../../{exp}/{dataset_name}/sample_{sample_id}_{label}.html">'
            f"{exp}</a></li>"
        )

    return f"""
    <div class="dropdown">
        <button class="btn btn-secondary dropdown-toggle" type="button"
            id="switchExperimentButton" data-bs-toggle="dropdown" aria-expanded="false">
            Switch probe
        </button>
        <ul class="dropdown-menu" aria-labelledby="switchExperimentButton">
            {experiment_links}
        </ul>
    </div>
    """


def get_full_html(
    title: str,
    navigation_buttons: list[str],
    content: str,
    metadata: str,
    subtitle: str | None = None,
) -> str:
    """Get full HTML page with content and metadata table.

    Args:
        title: Page title
        navigation_buttons: List of HTML strings for navigation buttons
        content: Main page content
        metadata: Dictionary of metadata to display in table
        subtitle: Optional subtitle to display below title (e.g. for score and label)
    """

    nav_buttons_html = '<div class="container mb-4"><div class="nav-buttons">'
    nav_buttons_html += "".join(navigation_buttons)
    nav_buttons_html += "</div></div>"

    title_html = f'<h1 class="display-4 mb-2">{title}</h1>'
    if subtitle:
        title_html += f'<h2 class="h4 mb-4">{subtitle}</h2>'
    else:
        title_html = title_html.replace("mb-2", "mb-4")

    full_html = 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>
        <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/css/bootstrap.min.css"
            rel="stylesheet">
        <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/js/bootstrap.bundle.min.js">
        </script>
        <link rel="stylesheet" href="../../style.css">

    </head>
    <body>
        <div class="container">
            {nav_buttons_html}
            {title_html}
            {content}
            {metadata}
        </div>
    </body>
    </html>
    """
    return full_html


def get_alpaca_idx_subset(df: pd.DataFrame) -> list[int]:
    """Get indices of most relevant samples from alpaca__plain."""
    lowest_scores = df.nsmallest(50, "score").index.tolist()
    first_500 = df.head(500).index.tolist()
    highest_scores = df.nlargest(50, "score").index.tolist()
    return list(set(lowest_scores + first_500 + highest_scores))  # type: ignore


def make_metadata_table(df_row: pd.Series) -> str:  # type: ignore
    """Make metadata table from dataframe row."""
    metadata_table = """
    <div class="container mt-4">
        <div class="d-flex align-items-center mb-3">
            <h3 class="h4 mb-0">Metadata</h3>
            <button class="btn btn-sm btn-secondary ms-2" type="button"
            data-bs-toggle="collapse" data-bs-target="#metadataTable"
            aria-expanded="false" aria-controls="metadataTable">
                Show metadata
            </button>
        </div>
        <hr class="mb-3">
        <div class="collapse" id="metadataTable">
            <table class="table table-sm">
                <tbody>
    """

    for k, v in df_row.items():
        key = html.escape(str(k))
        value = html.escape(str(v))  # type: ignore
        metadata_table += f"""
                    <tr>
                        <th scope="row">{key}</th>
                        <td>{value}</td>
                    </tr>"""
    metadata_table += """
                </tbody>
            </table>
        </div>
    </div>
    """
    return metadata_table


def export_dataset_html(
    experiment: Experiment,
    output_dir: Path,
    probe_name: str,
    probe_names: list[str],
    include_black_box: bool = False,
    skip_control: bool = False,
) -> dict[str, dict[str, int]]:
    """Export HTML files with token scores for each dataset.

    Args:
        experiment: Experiment object
        output_dir: Directory to save HTML files
        probe_name: Name of the probe
        probe_names: List of all probe names (for cross linking)

    Returns:
        Dictionary mapping dataset names to the number of samples
    """
    scorer = create_prompt_scorer("mean", experiment, [0.01], default_fpr=0.01, control_dataset_name="alpaca_XYZ__plain")
    output_dir.mkdir(parents=True, exist_ok=True)

    counts: dict[str, dict[str, int]] = {}

    # Choose datasets to process based on skip_control flag
    datasets_to_process = experiment.eval_datasets
    if not skip_control:
        datasets_to_process = experiment.eval_datasets + experiment.control_datasets

    for dataset in datasets_to_process:
        if dataset.name in EXCLUDED_DATASETS:
            continue
        # Get dataset info
        tok_dataset = TokenizedDataset.from_dataset(dataset, experiment.tokenizer, detect_all=experiment.cfg.detect_all_tokens)
        all_token_scores = get_scores_from_name(experiment, dataset.name, all_tokens=True)
        normalized_scores_array = scorer.normalize_scores(all_token_scores)
        df = get_dataframe_of_scores_on_dataset(
            experiment,
            dataset.name,
            scorer,
            normalize_scores=True,
            include_black_box=include_black_box,
        )

        # Get indices for honest, deceptive and ambiguous samples
        honest_idxs: list[int] = df[df["label"] == "honest"].index.tolist()
        deceptive_idxs: list[int] = df[df["label"] == "deceptive"].index.tolist()
        ambiguous_idxs: list[int] = df[df["label"] == "ambiguous"].index.tolist()

        if dataset.name == "alpaca__plain" or dataset.name == "alpaca_XYZ__plain":
            honest_idxs = get_alpaca_idx_subset(df)

        samples_info: list[dict[str, Any]] = []

        # Generate HTML for each sample
        for label, indices in [
            ("honest", honest_idxs),
            ("ambiguous", ambiguous_idxs),
            ("deceptive", deceptive_idxs),
        ]:
            for i in tqdm(indices, desc=f"{probe_name}, {dataset.name}, {label}"):  # type: ignore
                # Get sample info
                normalized_scores: Float[torch.Tensor, " n_tokens"] = normalized_scores_array[i]
                detection_mask = tok_dataset.detection_mask
                assert detection_mask is not None
                masked_scores = normalized_scores * detection_mask[i][-len(normalized_scores) :]
                str_tokens: list[str] = tok_dataset.str_tokens[i]

                # Store sample info
                df_row = df.iloc[i]
                score: float = df_row["score"]
                index = i  # Always use numeric index to avoid corrupted metadata
                sample_info: dict[str, int | str | float] = {
                    "id": index,
                    "label": str(label),
                    "score": score,
                }
                samples_info.append(sample_info)

                # Generate HTML for both masked and unmasked scores
                masked_html = get_token_scores_html(str_tokens, masked_scores, extent=6.0)
                unmasked_html = get_token_scores_html(str_tokens, normalized_scores, extent=6.0)

                # Add checkbox that toggles between masked and unmasked scores
                toggled_html = get_toggled_html(masked_html, unmasked_html)

                # Save HTML file

                navigation_buttons = [
                    "<a href='../../index.html' class='btn btn-secondary'>"
                    "    Go back to results overview"
                    "</a>",
                    "<a href='./index.html' class='btn btn-secondary'>"
                    "    Go back to dataset overview"
                    "</a>",
                    get_experiment_switch_button(
                        index, label, dataset.name, probe_names  # index is already int
                    ),
                ]

                metadata_table = make_metadata_table(df_row)  # type: ignore

                readable_name = to_readable_dataset_name_html(dataset.name)
                title = f"{readable_name} sample {index}"
                subtitle = f"Score: {score:.2f}, Label: {get_colored_label(label)}"

                full_html = get_full_html(
                    title=title,
                    navigation_buttons=navigation_buttons,
                    content=toggled_html,
                    metadata=metadata_table,
                    subtitle=subtitle,
                )
                filename = f"sample_{index}_{label}.html"
                dataset_output_dir = output_dir / dataset.name
                dataset_output_dir.mkdir(parents=True, exist_ok=True)
                with open(dataset_output_dir / filename, "w") as f:
                    f.write(full_html)

        # Create dataset-specific index
        counts[dataset.name] = create_dataset_index(
            dataset.name, samples_info, output_dir, probe_name, full_dataset_size=len(df)
        )

    return counts


class ProbeInfo(NamedTuple):
    name: str
    readable_name: str
    results_path: Path
    include_black_box: bool


def find_latest_experiment(results_path: Path) -> Path:
    """Find the latest experiment directory based on modification time."""
    experiment_dirs = [d for d in results_path.iterdir() if d.is_dir()]
    if not experiment_dirs:
        raise ValueError(f"No experiment directories found in {results_path}")
    
    # Sort by modification time (most recent first)
    latest_dir = max(experiment_dirs, key=lambda d: d.stat().st_mtime)
    return latest_dir


def generate(base_output_dir: Path | str = REPO_ROOT / "html-outputs", process_all: bool = True, skip_control: bool = False):
    RESULTS_PATH = REPO_ROOT / "results"
    base_output_dir = Path(base_output_dir)
    base_output_dir.mkdir(parents=True, exist_ok=True)
    
    if process_all:
        # Find ALL experiment directories, not just the latest
        experiment_dirs = [d for d in RESULTS_PATH.iterdir() if d.is_dir()]
        if not experiment_dirs:
            raise ValueError(f"No experiment directories found in {RESULTS_PATH}")
        
        # Sort by modification time (most recent first) for better organization
        experiment_dirs.sort(key=lambda d: d.stat().st_mtime, reverse=True)
        
        print(f"Found {len(experiment_dirs)} experiments to process:")
        for exp_dir in experiment_dirs:
            print(f"  - {exp_dir.name}")
        
        # Create probe info for all experiments
        probes = []
        for exp_dir in experiment_dirs:
            experiment_name = exp_dir.name
            probes.append(ProbeInfo(
                name=experiment_name,
                readable_name=f"Experiment ({experiment_name})",
                results_path=exp_dir,
                include_black_box=False,
            ))
    else:
        # Find the latest experiment automatically (original behavior)
        latest_experiment_path = find_latest_experiment(RESULTS_PATH)
        experiment_name = latest_experiment_path.name
        
        print(f"Using latest experiment: {experiment_name}")
        
        probes = [
            ProbeInfo(
                name=experiment_name,
                readable_name=f"Latest Experiment ({experiment_name})",
                results_path=latest_experiment_path,
                include_black_box=False,
            )
        ]
    # Write CSS file once in the base output directory
    write_style_css(base_output_dir)

    # Collect datasets for each probe and create HTML in each experiment folder
    all_experiment_info: dict[str, dict[str, Any]] = {}
    probe_names = [probe.name for probe in probes]

    for probe in probes:
        print(f"\n🔄 Processing experiment: {probe.name}")
        
        if process_all:
            # Create HTML output directory within the experiment folder
            html_output_dir = probe.results_path / "html"
            html_output_dir.mkdir(parents=True, exist_ok=True)
            
            # Write CSS file in each experiment's html folder
            write_style_css(html_output_dir)
        else:
            # Use the base output directory for single experiment
            html_output_dir = base_output_dir / probe.name

        print("probe.results_path", probe.results_path)

        # HACK: experiment.pkl should only exist if an upper_bound probe was trained
        experiment_pkl_path = probe.results_path / 'experiment.pkl'
        if experiment_pkl_path.exists():
            print("load experiment from ", experiment_pkl_path)
            with open(experiment_pkl_path, 'rb') as inp:
                experiment = pickle.load(inp)
        else:
            print(f"Creating lightweight experiment for {probe.name} (no model loading)")
            # For HTML export, we don't need the full model, just the scores
            experiment = Experiment(ExperimentConfig.from_path(probe.results_path))

        # Export HTML for this experiment
        counts = export_dataset_html(
            experiment, html_output_dir, probe.name, probe_names, probe.include_black_box, skip_control
        )
        
        # Store info for main index
        all_experiment_info[probe.name] = {
            "path": html_output_dir,
            "counts": counts,
            "readable_name": probe.readable_name
        }
        
        print(f"✅ Completed: {probe.name}")
        print(f"📁 HTML saved to: {html_output_dir.absolute()}")
        print(f"🌐 Open: {(html_output_dir / 'index.html').absolute()}")
            
    # Create main index page with all probes and datasets (only if not processing all)
    if not process_all:
        # Only include successfully processed experiments
        successful_experiments = {probe.name: all_experiment_info[probe.name]["counts"] 
                                 for probe in probes if probe.name in all_experiment_info}
        if successful_experiments:
            create_index_html(base_output_dir, successful_experiments)
        else:
            print("⚠️  No experiments were successfully processed for HTML export")
    
    print(f"\nHTML export completed!")
    print(f"📁 Files saved to: {base_output_dir.absolute()}")
    print(f"🌐 Open in browser: {(base_output_dir / 'index.html').absolute()}")


if __name__ == "__main__":
    fire.Fire(generate)
