from datetime import datetime
from pathlib import Path

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from rich.console import Console
from rich.table import Table
from rich.text import Text

from benchmark.config import PLOTS_DIR, RESULTS_DIR, AggregatedMetrics, EvalResults

matplotlib.use("Agg")


def print_summary_table(eval_results: EvalResults, k_values: list[int] | None = None):
    _print_rich_table(eval_results, k_values)


def _print_rich_table(eval_results: EvalResults, k_values: list[int] | None = None):

    console = Console()
    k_values = k_values or [1, 3, 5, 10, 20]

    table = Table(title="Retrieval Benchmark Results (File-Level)", show_lines=True)
    table.add_column("Retriever", style="bold cyan", min_width=18)
    for k in k_values:
        table.add_column(f"Recall@{k}", justify="center", min_width=9)
    table.add_column("MRR", justify="center", min_width=7)
    for k in [5, 10]:
        table.add_column(f"Hit@{k}", justify="center", min_width=7)
    table.add_column("Samples", justify="center", min_width=8)

    best = {}
    for metric_name in (
        [f"recall_{k}" for k in k_values] + ["mrr"] + [f"hit_{k}" for k in [5, 10]]
    ):
        best[metric_name] = 0.0

    for agg in eval_results.per_retriever:
        for k in k_values:
            best[f"recall_{k}"] = max(best[f"recall_{k}"], agg.recall_at_k.get(k, 0))
        best["mrr"] = max(best["mrr"], agg.mrr)
        for k in [5, 10]:
            best[f"hit_{k}"] = max(best[f"hit_{k}"], agg.hit_at_k.get(k, 0))

    for agg in eval_results.per_retriever:
        row = [agg.retriever]
        for k in k_values:
            val = agg.recall_at_k.get(k, 0)
            cell = f"{val:.3f}"
            if val == best[f"recall_{k}"] and val > 0:
                cell = f"[bold green]{cell}[/bold green]"
            row.append(cell)

        val = agg.mrr
        cell = f"{val:.3f}"
        if val == best["mrr"] and val > 0:
            cell = f"[bold green]{cell}[/bold green]"
        row.append(cell)

        for k in [5, 10]:
            val = agg.hit_at_k.get(k, 0)
            cell = f"{val:.3f}"
            if val == best[f"hit_{k}"] and val > 0:
                cell = f"[bold green]{cell}[/bold green]"
            row.append(cell)

        row.append(str(agg.num_samples))
        table.add_row(*row)

    console.print()
    console.print(table)

    has_method = any(
        any(v > 0 for v in agg.method_recall_at_k.values())
        for agg in eval_results.per_retriever
    )
    if has_method:
        mtable = Table(
            title="Retrieval Benchmark Results (Method-Level)", show_lines=True
        )
        mtable.add_column("Retriever", style="bold cyan", min_width=18)
        for k in k_values:
            mtable.add_column(f"Recall@{k}", justify="center", min_width=9)
        for k in [5, 10]:
            mtable.add_column(f"Hit@{k}", justify="center", min_width=7)

        for agg in eval_results.per_retriever:
            row = [agg.retriever]
            for k in k_values:
                row.append(f"{agg.method_recall_at_k.get(k, 0):.3f}")
            for k in [5, 10]:
                row.append(f"{agg.method_hit_at_k.get(k, 0):.3f}")
            mtable.add_row(*row)

        console.print()
        console.print(mtable)

    if eval_results.per_repo:
        console.print()
        console.print("[bold]Per-Repository Breakdown (Recall@5):[/bold]")

        repo_table = Table(show_lines=True)
        repo_table.add_column("Repository", style="bold", min_width=30)
        retrievers = [agg.retriever for agg in eval_results.per_retriever]
        for r in retrievers:
            repo_table.add_column(r, justify="center", min_width=10)
        repo_table.add_column("Samples", justify="center", min_width=8)

        for repo, aggs in sorted(eval_results.per_repo.items()):
            row = [repo]
            agg_by_name = {a.retriever: a for a in aggs}
            for r in retrievers:
                a = agg_by_name.get(r)
                row.append(f"{a.recall_at_k.get(5, 0):.3f}" if a else "-")
            row.append(str(aggs[0].num_samples if aggs else 0))
            repo_table.add_row(*row)

        console.print(repo_table)


def plot_results(eval_results: EvalResults, output_dir: Path | None = None):
    output_dir = output_dir or PLOTS_DIR
    output_dir.mkdir(parents=True, exist_ok=True)

    retrievers = [agg.retriever for agg in eval_results.per_retriever]
    colors = [
        "#2196F3",
        "#4CAF50",
        "#FF9800",
        "#9C27B0",
        "#F44336",
        "#00BCD4",
        "#795548",
    ]

    fig, ax = plt.subplots(figsize=(12, 6))
    metrics = {
        "Recall@1": [agg.recall_at_k.get(1, 0) for agg in eval_results.per_retriever],
        "Recall@5": [agg.recall_at_k.get(5, 0) for agg in eval_results.per_retriever],
        "Recall@10": [agg.recall_at_k.get(10, 0) for agg in eval_results.per_retriever],
        "MRR": [agg.mrr for agg in eval_results.per_retriever],
        "Hit@5": [agg.hit_at_k.get(5, 0) for agg in eval_results.per_retriever],
    }


    x = np.arange(len(metrics))
    width = 0.8 / len(retrievers)

    for i, (retriever, color) in enumerate(zip(retrievers, colors)):
        values = [metrics[m][i] for m in metrics]
        offset = (i - len(retrievers) / 2 + 0.5) * width
        ax.bar(x + offset, values, width, label=retriever, color=color, alpha=0.85)

    ax.set_xticks(x)
    ax.set_xticklabels(metrics.keys())
    ax.set_ylabel("Score")
    ax.set_title("Retrieval Benchmark: Metrics Comparison")
    ax.legend(loc="upper left", fontsize=8)
    ax.set_ylim(0, 1.05)
    ax.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_dir / "metrics_comparison.png", dpi=150)
    plt.close()

    fig, ax = plt.subplots(figsize=(10, 6))
    k_values = (
        sorted(eval_results.per_retriever[0].recall_at_k.keys())
        if eval_results.per_retriever
        else []
    )

    for i, (agg, color) in enumerate(zip(eval_results.per_retriever, colors)):
        recalls = [agg.recall_at_k.get(k, 0) for k in k_values]
        ax.plot(
            k_values,
            recalls,
            marker="o",
            color=color,
            label=agg.retriever,
            linewidth=2,
            markersize=6,
        )

    ax.set_xlabel("K (top-K)")
    ax.set_ylabel("Recall@K")
    ax.set_title("Recall@K Curve by Retriever")
    ax.legend(fontsize=9)
    ax.set_ylim(0, 1.05)
    ax.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_dir / "recall_at_k_curve.png", dpi=150)
    plt.close()

    if eval_results.per_repo and len(eval_results.per_repo) > 1:
        repos = sorted(eval_results.per_repo.keys())
        data = []
        for repo in repos:
            aggs = {a.retriever: a for a in eval_results.per_repo[repo]}
            row = [
                aggs[r].recall_at_k.get(5, 0) if r in aggs else 0 for r in retrievers
            ]
            data.append(row)

        fig, ax = plt.subplots(
            figsize=(max(8, len(retrievers) * 2), max(4, len(repos) * 0.6))
        )
        im = ax.imshow(data, cmap="YlOrRd", aspect="auto", vmin=0, vmax=1)

        ax.set_xticks(range(len(retrievers)))
        ax.set_xticklabels(retrievers, rotation=45, ha="right", fontsize=9)
        ax.set_yticks(range(len(repos)))
        short_repos = [r.split("/")[-1] if "/" in r else r for r in repos]
        ax.set_yticklabels(short_repos, fontsize=9)

        for i in range(len(repos)):
            for j in range(len(retrievers)):
                ax.text(
                    j,
                    i,
                    f"{data[i][j]:.2f}",
                    ha="center",
                    va="center",
                    fontsize=8,
                    color="white" if data[i][j] > 0.5 else "black",
                )

        plt.colorbar(im, label="Recall@5")
        ax.set_title("Recall@5 per Repository")
        plt.tight_layout()
        plt.savefig(output_dir / "per_repo_heatmap.png", dpi=150)
        plt.close()

    print(f"Plots saved to {output_dir}")


def save_results_markdown(
    eval_results: EvalResults,
    k_values: list[int] | None = None,
    output_path: Path | None = None,
):
    k_values = k_values or [1, 3, 5, 10, 20]
    output_path = output_path or (RESULTS_DIR / "benchmark_log.md")

    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    retrievers = [agg.retriever for agg in eval_results.per_retriever]
    n_samples = (
        eval_results.per_retriever[0].num_samples if eval_results.per_retriever else 0
    )

    lines = []
    lines.append(f"\n## Run: {timestamp}")
    lines.append(f"Samples: {n_samples} | Retrievers: {', '.join(retrievers)}\n")

    # File-level table
    lines.append("### File-Level Metrics")
    header = "| Retriever |"
    sep = "|-----------|"
    for k in k_values:
        header += f"R@{k} |"
        sep += "------|"
    header += " MRR | Hit@5 | Hit@10 |"
    sep += "------|-------|--------|"
    lines.append(header)
    lines.append(sep)

    for agg in eval_results.per_retriever:
        row = f"| {agg.retriever} |"
        for k in k_values:
            row += f"{agg.recall_at_k.get(k, 0):.3f} |"
        row += f"{agg.mrr:.3f} | {agg.hit_at_k.get(5, 0):.3f} | {agg.hit_at_k.get(10, 0):.3f} |"
        lines.append(row)

    # Per-repo breakdown
    if eval_results.per_repo:
        lines.append("\n### Per-Repository (Recall@5)")
        header = "| Repository |"
        sep = "|------------|"
        for r in retrievers:
            header += f"{r} |"
            sep += "------|"
        header += " N |"
        sep += "---|"
        lines.append(header)
        lines.append(sep)

        for repo, aggs in sorted(eval_results.per_repo.items()):
            agg_by_name = {a.retriever: a for a in aggs}
            row = f"| {repo} |"
            for r in retrievers:
                a = agg_by_name.get(r)
                row += f"{a.recall_at_k.get(5, 0):.3f} |" if a else " - |"
            row += f"{aggs[0].num_samples if aggs else 0} |"
            lines.append(row)

    lines.append("")

    if not output_path.exists():
        with open(output_path, "w") as f:
            f.write("# Retrieval Benchmark Results Log\n")

    with open(output_path, "a") as f:
        f.write("\n".join(lines) + "\n")

    print(f"Results appended to {output_path}")