# platform/deskclaw/experiment/reporter.py
# ═══════════════════════════════════════════════════════════════
# ComparativeReporter — Markdown tables and matplotlib charts
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

from deskclaw.experiment.models import ExperimentResult
from deskclaw.experiment.scoring import ExperimentMetrics

logger = logging.getLogger("deskclaw.experiment.reporter")


class ComparativeReporter:
    def __init__(
        self,
        result: ExperimentResult,
        metrics: ExperimentMetrics,
    ) -> None:
        self._result = result
        self._metrics = metrics

    def generate_markdown(self) -> str:
        lines: list[str] = []
        lines.append(f"# Experiment Report: {self._result.config.name}\n")
        lines.append(f"- **Experiment ID**: {self._result.experiment_id}")
        lines.append(f"- **Started**: {self._result.started_at}")
        lines.append(f"- **Finished**: {self._result.finished_at}")
        lines.append(f"- **Total Runs**: {self._metrics.total_runs}")
        lines.append(f"- **Overall Pass Rate**: {self._metrics.overall_pass_rate:.2%}")
        lines.append(f"- **Total Tokens**: {self._metrics.total_tokens:,}")
        lines.append(f"- **Total Cost**: ${self._metrics.total_cost_usd:.4f}")
        lines.append("")

        lines.append("## Blueprint Comparison\n")
        lines.append(
            "| Blueprint | Runs | Pass Rate | Mean Latency (ms) | P95 Latency (ms) "
            "| Total Tokens | Cost ($) |"
        )
        lines.append(
            "|-----------|-----:|----------:|-------------------:|------------------:"
            "|-------------:|---------:|"
        )
        for bp_name, bp in sorted(self._metrics.blueprint_metrics.items()):
            lines.append(
                f"| {bp_name} | {bp.run_count} | {bp.pass_rate:.2%} "
                f"| {bp.mean_latency_ms:.0f} | {bp.p95_latency_ms:.0f} "
                f"| {bp.total_tokens:,} | {bp.total_cost_usd:.4f} |"
            )
        lines.append("")

        lines.append("## Per-Blueprint Details\n")
        for bp_name, bp in sorted(self._metrics.blueprint_metrics.items()):
            lines.append(f"### {bp_name}\n")
            lines.append(f"- Runs: {bp.run_count}")
            lines.append(f"- Mean Latency: {bp.mean_latency_ms:.2f} ms")
            lines.append(f"- Median Latency: {bp.median_latency_ms:.2f} ms")
            lines.append(f"- P95 Latency: {bp.p95_latency_ms:.2f} ms")
            lines.append(f"- P99 Latency: {bp.p99_latency_ms:.2f} ms")
            lines.append(f"- Input Tokens: {bp.total_input_tokens:,}")
            lines.append(f"- Output Tokens: {bp.total_output_tokens:,}")
            lines.append(f"- Mean Tokens/Run: {bp.mean_tokens_per_run:.0f}")
            lines.append(f"- Total Cost: ${bp.total_cost_usd:.4f}")
            lines.append(f"- Mean Cost/Run: ${bp.mean_cost_per_run:.6f}")
            lines.append("")

        matrix = self._result.config.matrix.to_dict()
        if matrix:
            lines.append("## Variable Matrix\n")
            for var_name, values in sorted(matrix.items()):
                lines.append(f"- **{var_name}**: {values}")
            lines.append("")

        return "\n".join(lines)

    def generate_plots(self, output_dir: str | Path) -> list[Path]:
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        plots: list[Path] = []

        try:
            import matplotlib

            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError:
            logger.warning("matplotlib not available — skipping plot generation")
            return plots

        bp_names = sorted(self._metrics.blueprint_metrics.keys())
        if not bp_names:
            return plots

        plots.append(self._plot_latency(plt, bp_names, output_dir))
        plots.append(self._plot_tokens(plt, bp_names, output_dir))
        plots.append(self._plot_cost(plt, bp_names, output_dir))
        plt.close("all")

        return [p for p in plots if p is not None]

    def write_report(self, output_dir: str | Path, fmt: str = "both") -> list[Path]:
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        written: list[Path] = []

        if fmt in ("md", "both"):
            md_path = output_dir / "experiment_report.md"
            md_path.write_text(self.generate_markdown(), encoding="utf-8")
            written.append(md_path)
            logger.info("Wrote markdown report: %s", md_path)

        if fmt in ("json", "both"):
            json_path = output_dir / "experiment_report.json"
            import json

            data = {
                "experiment": self._result.to_dict(),
                "metrics": self._metrics.to_dict(),
            }
            json_path.write_text(
                json.dumps(data, indent=2, ensure_ascii=False, default=str),
                encoding="utf-8",
            )
            written.append(json_path)

        if fmt in ("plots", "both"):
            written.extend(self.generate_plots(output_dir))

        return written

    def _plot_latency(
        self, plt: Any, bp_names: list[str], output_dir: Path
    ) -> Path | None:
        try:
            metrics = self._metrics.blueprint_metrics
            means = [metrics[n].mean_latency_ms for n in bp_names]
            p95s = [metrics[n].p95_latency_ms for n in bp_names]

            fig, ax = plt.subplots(figsize=(max(8, len(bp_names) * 2), 5))
            x = range(len(bp_names))
            width = 0.35
            ax.bar([i - width / 2 for i in x], means, width, label="Mean")
            ax.bar([i + width / 2 for i in x], p95s, width, label="P95")
            ax.set_xlabel("Blueprint")
            ax.set_ylabel("Latency (ms)")
            ax.set_title("Latency Comparison")
            ax.set_xticks(list(x))
            ax.set_xticklabels(bp_names, rotation=45, ha="right")
            ax.legend()
            fig.tight_layout()

            path = output_dir / "latency_comparison.png"
            fig.savefig(str(path), dpi=100)
            return path
        except Exception as exc:
            logger.warning("Failed to generate latency plot: %s", exc)
            return None

    def _plot_tokens(
        self, plt: Any, bp_names: list[str], output_dir: Path
    ) -> Path | None:
        try:
            metrics = self._metrics.blueprint_metrics
            inputs = [metrics[n].total_input_tokens for n in bp_names]
            outputs = [metrics[n].total_output_tokens for n in bp_names]

            fig, ax = plt.subplots(figsize=(max(8, len(bp_names) * 2), 5))
            x = range(len(bp_names))
            width = 0.35
            ax.bar([i - width / 2 for i in x], inputs, width, label="Input")
            ax.bar([i + width / 2 for i in x], outputs, width, label="Output")
            ax.set_xlabel("Blueprint")
            ax.set_ylabel("Tokens")
            ax.set_title("Token Usage Comparison")
            ax.set_xticks(list(x))
            ax.set_xticklabels(bp_names, rotation=45, ha="right")
            ax.legend()
            fig.tight_layout()

            path = output_dir / "token_usage.png"
            fig.savefig(str(path), dpi=100)
            return path
        except Exception as exc:
            logger.warning("Failed to generate token plot: %s", exc)
            return None

    def _plot_cost(
        self, plt: Any, bp_names: list[str], output_dir: Path
    ) -> Path | None:
        try:
            metrics = self._metrics.blueprint_metrics
            costs = [metrics[n].total_cost_usd for n in bp_names]

            fig, ax = plt.subplots(figsize=(max(8, len(bp_names) * 2), 5))
            ax.bar(bp_names, costs)
            ax.set_xlabel("Blueprint")
            ax.set_ylabel("Cost (USD)")
            ax.set_title("Cost Comparison")
            plt.xticks(rotation=45, ha="right")
            fig.tight_layout()

            path = output_dir / "cost_comparison.png"
            fig.savefig(str(path), dpi=100)
            return path
        except Exception as exc:
            logger.warning("Failed to generate cost plot: %s", exc)
            return None
