import html
import json
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional

from tabulate import tabulate


@dataclass
class OptimizationResult:
    """Represents the optimization results for a single story."""

    cleaned_outputs: List[str]
    targets: List[float]
    xentropies: List[float]
    new_predictions: List[str]
    success_score: int

    @property
    def best_context_index(self) -> Optional[int]:
        """Returns the index of the best optimized context based on target score."""
        if not self.targets:
            return None
        return max(range(len(self.targets)), key=lambda i: self.targets[i])

    @property
    def best_context(self) -> Optional[str]:
        """Returns the best optimized context based on target score."""
        idx = self.best_context_index
        if idx is None:
            return None
        return self.cleaned_outputs[idx]

    @property
    def best_prediction(self) -> Optional[str]:
        """Returns the prediction corresponding to the best context."""
        idx = self.best_context_index
        if idx is None:
            return None
        return self.new_predictions[idx]


@dataclass
class StoryExperiment:
    """Represents a single story experiment with a template, default context, and optimization results."""

    story_template: str
    default_context: str
    desired_word: str
    undesired_word: str
    optimization_results: OptimizationResult

    @property
    def is_successful(self) -> bool:
        """Returns whether the optimization was successful."""
        return self.optimization_results.success_score > 0

    def get_full_story_with_default_context(self) -> str:
        """Returns the full story with the default context."""
        return self.story_template.format(self.default_context)

    def get_full_story_with_best_context(self) -> str:
        """Returns the full story with the best optimized context."""
        best_context = self.optimization_results.best_context
        if best_context is None:
            return self.get_full_story_with_default_context()
        return self.story_template.format(best_context)


class StoryDataset:
    """Collection of story experiments loaded from JSON."""

    def __init__(self, experiments: List[StoryExperiment]):
        self.experiments = experiments

    @classmethod
    def from_json_file(cls, json_path: str) -> "StoryDataset":
        """Load story experiments from a JSON file."""
        path = Path(json_path)
        if not path.exists():
            raise FileNotFoundError(f"JSON file not found: {json_path}")

        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)

        experiments = []
        for item in data:
            opt_results = item["optimization_results"]

            # Create OptimizationResult object
            optimization_result = OptimizationResult(
                cleaned_outputs=opt_results["cleaned_outputs"],
                targets=opt_results["targets"],
                xentropies=opt_results["xentropies"],
                new_predictions=opt_results["new_predictions"],
                success_score=opt_results["success_score"],
            )

            # Create StoryExperiment object
            experiment = StoryExperiment(
                story_template=item["story_template"],
                default_context=item["default_context"],
                desired_word=item["desired_word"],
                undesired_word=item["undesired_word"],
                optimization_results=optimization_result,
            )

            experiments.append(experiment)

        return cls(experiments)

    @property
    def successful_experiments(self) -> List[StoryExperiment]:
        """Returns a list of successful experiments."""
        return [exp for exp in self.experiments if exp.is_successful]

    @property
    def failed_experiments(self) -> List[StoryExperiment]:
        """Returns a list of failed experiments."""
        return [exp for exp in self.experiments if not exp.is_successful]

    @property
    def success_rate(self) -> float:
        """Returns the percentage of successful experiments."""
        if not self.experiments:
            return 0.0
        return len(self.successful_experiments) / len(self.experiments)

    def summary(self) -> str:
        """Returns a summary of the dataset."""
        return (
            f"Story Experiments: {len(self.experiments)}\n"
            f"Successful: {len(self.successful_experiments)} ({self.success_rate:.1%})\n"
            f"Failed: {len(self.failed_experiments)} ({1-self.success_rate:.1%})"
        )

    def generate_html_report(self, output_path: str) -> None:
        """
        Generate an HTML report with tabs for each experiment and write it to a file.

        Args:
            output_path: Path to save the HTML report
        """
        # Start HTML document with Bootstrap for styling
        html_content = """<!DOCTYPE html>
        <html lang="en">
        <head>
            <meta charset="UTF-8">
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <title>Story Experiments Report</title>
            <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0-alpha1/dist/css/bootstrap.min.css" rel="stylesheet">
            <style>
                body { padding: 20px; }
                .tab-content { padding: 20px 0; }
                pre { white-space: pre-wrap; background-color: #f8f9fa; padding: 10px; border-radius: 4px; }
                .context-box { background-color: #f0fff0; padding: 15px; border-radius: 5px; margin: 10px 0; }
                .nav-tabs { margin-bottom: 20px; }
                .table-responsive { margin-top: 20px; }
                .metrics { font-weight: bold; }
                .best-row { background-color: #e6ffe6; }
            </style>
        </head>
        <body>
            <div class="container-fluid">
                <h1>Story Experiments Report</h1>
                <div class="row">
                    <div class="col-12">
                        <ul class="nav nav-tabs" id="experimentTabs" role="tablist">
                            <li class="nav-item" role="presentation">
                                <button class="nav-link active" id="summary-tab" data-bs-toggle="tab" data-bs-target="#summary"
                                        type="button" role="tab" aria-controls="summary" aria-selected="true">Summary</button>
                            </li>
        """

        # Add tab buttons for each experiment
        for i, exp in enumerate(self.experiments):
            exp_id = f"exp-{i}"
            active = ""
            html_content += f"""
                            <li class="nav-item" role="presentation">
                                <button class="nav-link {active}" id="{exp_id}-tab" data-bs-toggle="tab" data-bs-target="#{exp_id}"
                                        type="button" role="tab" aria-controls="{exp_id}" aria-selected="false">Exp {i+1}</button>
                            </li>"""

        html_content += """
                        </ul>

                        <div class="tab-content" id="experimentTabContent">
                            <div class="tab-pane fade show active" id="summary" role="tabpanel" aria-labelledby="summary-tab">
                                <h2>Summary</h2>
                                <pre>{}</pre>
        """.format(
            html.escape(self.summary())
        )

        # Add a table for successful experiments in the summary tab
        html_content += """
                                <h3>Successful Experiments Overview</h3>
                                <div class="table-responsive">
                                    <table class="table table-striped table-bordered">
                                        <thead>
                                            <tr>
                                                <th>Exp #</th>
                                                <th>Desired Word</th>
                                                <th>Undesired Word</th>
                                                <th>Target</th>
                                                <th>XEntropy</th>
                                                <th>Success</th>
                                            </tr>
                                        </thead>
                                        <tbody>
        """

        for i, exp in enumerate(self.experiments):
            best_idx = exp.optimization_results.best_context_index
            if best_idx is not None:
                target = exp.optimization_results.targets[best_idx]
                xentropy = exp.optimization_results.xentropies[best_idx]
                success = "✅" if exp.is_successful else "❌"

                html_content += f"""
                                            <tr>
                                                <td>{i+1}</td>
                                                <td>{html.escape(exp.desired_word)}</td>
                                                <td>{html.escape(exp.undesired_word)}</td>
                                                <td>{target:.4f}</td>
                                                <td>{xentropy:.4f}</td>
                                                <td>{success}</td>
                                            </tr>
                """

        html_content += """
                                        </tbody>
                                    </table>
                                </div>
                            </div>
        """

        # Add tab content for each experiment
        for i, exp in enumerate(self.experiments):
            exp_id = f"exp-{i}"
            active = ""
            best_idx = exp.optimization_results.best_context_index

            success_text = "Successful" if exp.is_successful else "Failed"
            success_class = "text-success" if exp.is_successful else "text-danger"

            html_content += f"""
                            <div class="tab-pane fade {active}" id="{exp_id}" role="tabpanel" aria-labelledby="{exp_id}-tab">
                                <h2>Experiment {i+1} <span class="{success_class}">({success_text})</span></h2>

                                <div class="row">
                                    <div class="col-12">
                                        <h3>Experiment Details</h3>
                                        <p><strong>Desired Word:</strong> {html.escape(exp.desired_word)}</p>
                                        <p><strong>Undesired Word:</strong> {html.escape(exp.undesired_word)}</p>
                                        <p><strong>Story Template:</strong></p>
                                        <pre>{html.escape(exp.story_template)}</pre>
                                        <p><strong>Default Context:</strong></p>
                                        <div class="context-box">{html.escape(exp.default_context)}</div>
                                    </div>
                                </div>

                                <h3>All Optimization Attempts</h3>
                                <div class="table-responsive">
                                    <table class="table table-striped table-bordered">
                                        <thead>
                                            <tr>
                                                <th>#</th>
                                                <th>Context</th>
                                                <th>Prediction</th>
                                                <th>Target Score</th>
                                                <th>Cross-Entropy</th>
                                            </tr>
                                        </thead>
                                        <tbody>
            """

            # Add rows for all optimization attempts
            for j, (context, target, xentropy, prediction) in enumerate(
                zip(
                    exp.optimization_results.cleaned_outputs,
                    exp.optimization_results.targets,
                    exp.optimization_results.xentropies,
                    exp.optimization_results.new_predictions,
                )
            ):
                row_class = "best-row" if j == best_idx else ""
                best_label = " (Best)" if j == best_idx else ""

                html_content += f"""
                                            <tr class="{row_class}">
                                                <td>{j+1}{best_label}</td>
                                                <td><pre style="margin: 0">{html.escape(context)}</pre></td>
                                                <td><pre style="margin: 0">{html.escape(prediction)}</pre></td>
                                                <td>{target:.4f}</td>
                                                <td>{xentropy:.4f}</td>
                                            </tr>
                """

            html_content += """
                                        </tbody>
                                    </table>
                                </div>
                            </div>
            """

        # Close the HTML document
        html_content += """
                        </div>
                    </div>
                </div>
            </div>

            <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0-alpha1/dist/js/bootstrap.bundle.min.js"></script>
        </body>
        </html>
        """

        # Write the HTML content to a file
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(html_content)

        print(f"HTML report generated and saved to: {output_path}")


# Example usage
if __name__ == "__main__":
    dataset = StoryDataset.from_json_file("optimization_results_01.json")
    print(dataset.summary())

    # Print details of successful experiments in a table
    print("\nSuccessful Experiments:")

    table_data = []
    for i, exp in enumerate(dataset.successful_experiments):
        best_idx = exp.optimization_results.best_context_index
        if best_idx is not None:
            target = exp.optimization_results.targets[best_idx]
            xentropy = exp.optimization_results.xentropies[best_idx]

            table_data.append(
                [
                    i + 1,
                    exp.desired_word,
                    exp.undesired_word,
                    exp.default_context,
                    exp.optimization_results.best_context,
                    exp.optimization_results.best_prediction,
                    f"{target:.4f}",
                    f"{xentropy:.4f}",
                ]
            )

    headers = [
        "Exp #",
        "Desired Word",
        "Undesired Word",
        "Default Context",
        "Best Context",
        "Prediction",
        "Target",
        "XEntropy",
    ]

    print(tabulate(table_data, headers=headers, tablefmt="grid"))

    # Generate HTML report
    dataset.generate_html_report("view2.html")
