"""Reporting and visualization for ablation runs."""

from __future__ import annotations

import json
import math
import re
import warnings
from dataclasses import dataclass, field
from pathlib import Path, PureWindowsPath
from typing import Any

import pandas as pd

REPORT_SCHEMA_VERSION = 1
_SENSITIVE_KEY_RE = re.compile(
    r"(?:authorization|credential|password|secret|token|api[_-]?key)", re.IGNORECASE,
)


def _sanitize_label(text: str, max_len: int = 80) -> str:
    """Strip filesystem paths, tokens, and overly-long strings from labels."""
    if text.startswith("/"):
        text = Path(text).name
    elif re.match(r"^[A-Za-z]:[\\/]", text):
        text = PureWindowsPath(text).name
    text = re.sub(
        r"(?:/[A-Za-z0-9_.-]+){2,}", lambda match: Path(match.group()).name, text,
    )
    text = re.sub(r"\bhf_[A-Za-z0-9]{6,}\b", "<TOKEN>", text)
    text = re.sub(r"\bgh[pousr]_[A-Za-z0-9]{12,}\b", "<TOKEN>", text)
    text = re.sub(r"\bgithub_pat_[A-Za-z0-9_]{12,}\b", "<TOKEN>", text)
    text = re.sub(r"\bsk-[A-Za-z0-9_-]{12,}\b", "<TOKEN>", text)
    text = re.sub(r"\b[0-9a-fA-F]{32,}\b", "<REDACTED>", text)
    if len(text) > max_len:
        text = text[: max_len - 3] + "..."
    return text


def _sanitize_public_value(value: Any) -> Any:
    """Return a deterministic, JSON-safe value with private material removed."""
    if value is None or isinstance(value, (bool, int)):
        return value
    if isinstance(value, float):
        return value if math.isfinite(value) else None
    if isinstance(value, str):
        return _sanitize_label(value, max_len=240)
    if isinstance(value, dict):
        return {
            str(key): _sanitize_public_value(item)
            for key, item in sorted(value.items(), key=lambda pair: str(pair[0]))
            if not _SENSITIVE_KEY_RE.search(str(key))
        }
    if isinstance(value, (list, tuple)):
        return [_sanitize_public_value(item) for item in value]
    return _sanitize_label(str(value), max_len=240)


def _normalize_metrics(metrics: dict[str, Any]) -> tuple[dict[str, float | None], dict[str, str]]:
    values: dict[str, float | None] = {}
    status: dict[str, str] = {}
    for name, raw_value in sorted(metrics.items()):
        value = None
        if not isinstance(raw_value, bool):
            try:
                candidate = float(raw_value)
                if math.isfinite(candidate):
                    value = candidate
            except (TypeError, ValueError):
                pass
        values[str(name)] = value
        status[str(name)] = "measured" if value is not None else "unavailable"
    return values, status


@dataclass
class AblationResult:
    """Result of a single ablation experiment."""

    strategy: str
    component: str
    description: str
    metrics: dict[str, float | None]
    metadata: dict[str, Any] | None = None


@dataclass
class AblationReport:
    """Collects results and produces tables / charts / exports."""

    model_name: str
    baseline_metrics: dict[str, float | None] = field(default_factory=dict)
    results: list[AblationResult] = field(default_factory=list)

    def add_baseline(self, metrics: dict[str, float | None]):
        self.baseline_metrics = metrics

    def add_result(self, result: AblationResult):
        self.results.append(result)

    def to_dict(self) -> dict[str, Any]:
        """Return the canonical, versioned public report representation."""
        baseline_metrics, baseline_status = _normalize_metrics(self.baseline_metrics)
        results = []
        for result in self.results:
            metrics, metric_status = _normalize_metrics(result.metrics)
            results.append({
                "strategy": _sanitize_label(result.strategy),
                "component": _sanitize_label(result.component),
                "description": _sanitize_label(result.description, max_len=240),
                "metrics": metrics,
                "metric_status": metric_status,
                "metadata": _sanitize_public_value(result.metadata),
            })
        return {
            "schema_version": REPORT_SCHEMA_VERSION,
            "model_name": _sanitize_label(self.model_name),
            "baseline_metrics": baseline_metrics,
            "baseline_metric_status": baseline_status,
            "results": results,
        }

    def to_dataframe(self) -> pd.DataFrame:
        """Convert results to a pandas DataFrame with delta columns."""
        rows = []
        baseline_metrics, _ = _normalize_metrics(self.baseline_metrics)
        for r in self.results:
            metrics, _ = _normalize_metrics(r.metrics)
            row = {
                "strategy": _sanitize_label(r.strategy),
                "component": _sanitize_label(r.component),
                "description": _sanitize_label(r.description, max_len=240),
            }
            for metric_name, value in metrics.items():
                row[metric_name] = value
                baseline_val = baseline_metrics.get(metric_name)
                if baseline_val is not None and value is not None:
                    row[f"{metric_name}_delta"] = value - baseline_val
                    if baseline_val != 0:
                        row[f"{metric_name}_pct_change"] = (
                            (value - baseline_val) / abs(baseline_val)
                        ) * 100
            rows.append(row)

        return pd.DataFrame(rows)

    def print_summary(self):
        """Print a rich-formatted summary table."""
        from rich.console import Console
        from rich.table import Table

        console = Console()
        df = self.to_dataframe()

        if df.empty:
            console.print("[yellow]No ablation results to display.[/yellow]")
            return

        table = Table(title=f"Ablation Results: {_sanitize_label(self.model_name)}")
        table.add_column("Strategy", style="cyan")
        table.add_column("Component", style="green")

        baseline_metrics, _ = _normalize_metrics(self.baseline_metrics)
        metric_names = list(baseline_metrics.keys())
        for m in metric_names:
            table.add_column(f"{m}", justify="right")
            table.add_column(f"{m} delta", justify="right", style="red")

        # Baseline row
        baseline_vals = []
        for m in metric_names:
            value = baseline_metrics[m]
            baseline_vals.extend([f"{value:.4f}" if value is not None else "unavailable", "—"])
        table.add_row("baseline", "—", *baseline_vals, style="bold")

        for _, row in df.iterrows():
            cells = [row["strategy"], row["component"]]
            for m in metric_names:
                val = row.get(m, float("nan"))
                delta = row.get(f"{m}_delta", float("nan"))
                cells.append(f"{val:.4f}" if pd.notna(val) else "unavailable")
                cells.append(f"{delta:+.4f}" if pd.notna(delta) else "—")
            table.add_row(*cells)

        console.print(table)

    def save_json(self, path: str | Path):
        """Save raw results to JSON."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(
            json.dumps(self.to_dict(), indent=2, sort_keys=True, allow_nan=False) + "\n",
            encoding="utf-8",
        )

    def save_csv(self, path: str | Path):
        """Save results DataFrame to CSV."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        self.to_dataframe().to_csv(path, index=False, na_rep="unavailable", lineterminator="\n")

    def plot_impact(self, metric: str | None = None, output_path: str | Path | None = None):
        """Generate a bar chart showing the impact of each ablation on a metric.

        Args:
            metric: Which metric to plot. Defaults to the first baseline metric.
            output_path: If provided, save the figure instead of showing it.
        """
        import matplotlib

        if output_path:
            matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import seaborn as sns

        if metric is None:
            metric = list(self.baseline_metrics.keys())[0]

        df = self.to_dataframe()
        delta_col = f"{metric}_delta"
        if delta_col not in df.columns:
            raise ValueError(f"No delta column for metric {metric!r}")

        df_sorted = df.sort_values(delta_col, ascending=True)

        fig, ax = plt.subplots(figsize=(12, max(4, len(df_sorted) * 0.35)))
        colors = ["#e74c3c" if v > 0 else "#2ecc71" for v in df_sorted[delta_col]]
        sns.barplot(
            x=delta_col, y="component", hue="component", data=df_sorted,
            palette=dict(zip(df_sorted["component"], colors)), legend=False, ax=ax,
        )

        ax.set_xlabel(f"Change in {metric} (vs baseline)")
        ax.set_ylabel("Ablated Component")
        ax.set_title(f"Ablation Impact on {metric} — {_sanitize_label(self.model_name)}")
        ax.axvline(x=0, color="black", linewidth=0.8)

        plt.tight_layout()
        if output_path:
            fig.savefig(output_path, dpi=150, bbox_inches="tight")
            plt.close(fig)
        else:
            plt.show()

    def plot_heatmap(self, output_path: str | Path | None = None):
        """Generate a heatmap of pct_change across all strategies and metrics."""
        import matplotlib

        if output_path:
            matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import seaborn as sns

        df = self.to_dataframe()
        pct_cols = [c for c in df.columns if c.endswith("_pct_change")]
        if not pct_cols:
            return

        pivot = df.set_index("component")[pct_cols]
        pivot.columns = [c.replace("_pct_change", "") for c in pivot.columns]

        fig, ax = plt.subplots(figsize=(max(6, len(pivot.columns) * 2), max(4, len(pivot) * 0.4)))
        # Seaborn currently calls Matplotlib's pending-deprecated ``set_bad``
        # internally. Keep the repository's zero-warning contract focused on
        # our code while containing this specific third-party compatibility
        # warning to the call that emits it.
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="The set_bad function will be deprecated.*",
                category=PendingDeprecationWarning,
            )
            sns.heatmap(pivot, annot=True, fmt=".1f", cmap="RdYlGn_r", center=0, ax=ax)
        ax.set_title(f"Ablation % Change — {_sanitize_label(self.model_name)}")

        plt.tight_layout()
        if output_path:
            fig.savefig(output_path, dpi=150, bbox_inches="tight")
            plt.close(fig)
        else:
            plt.show()