"""Benchmark runner for executing evaluations."""

import csv
import json
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, cast

import lancedb
from benchmarks.datasets import HuggingFaceDataset, LocalDataset
from benchmarks.tasks import get_task
from configs.models import BenchmarkConfig, DatasetSourceConfig, ModelConfig
from rich.console import Console
from rich.table import Table

console = Console()


def _create_task(config: BenchmarkConfig, model_config: ModelConfig | None) -> Any:
    """Create benchmark task with task-specific model settings.

    Args:
        config: Benchmark configuration.
        model_config: Optional model configuration for task-owned loading.

    Returns:
        Benchmark task instance.
    """
    task_kwargs: dict[str, Any] = {"top_k_list": config.task.top_k_list}

    if config.task.type == "retrieval" and model_config is not None:
        task_kwargs.update(
            {
                "dino_model": model_config.dino_model,
                "compression_dim": model_config.compression_dim,
                "compressor_path": model_config.compressor_path,
            }
        )

    return get_task(config.task.type, **task_kwargs)


def create_dataset(config: DatasetSourceConfig) -> Any:
    """Create a dataset instance from configuration.

    Args:
        config: Dataset source configuration.

    Returns:
        Dataset instance.

    Raises:
        ValueError: If source_type is not supported.
    """
    if config.source_type == "huggingface":
        return HuggingFaceDataset(
            hf_id=config.path,
            img_column=config.img_column,
            label_column=config.label_column,
        )
    elif config.source_type == "local":
        return LocalDataset(
            local_path=config.path,
            img_column=config.img_column,
            label_column=config.label_column,
        )
    else:
        raise ValueError(
            f"Unsupported source_type: {config.source_type}. "
            f"Supported types: 'huggingface', 'local'"
        )


def _get_table_name(
    config: BenchmarkConfig,
    model_name: str,
    dataset_path: str,
) -> str:
    """Generate database table name from config, model, and dataset.

    Args:
        config: Benchmark configuration.
        model_name: Model name for table naming.
        dataset_path: Dataset identifier or path (e.g. "uoft-cs/cifar100").

    Returns:
        Formatted table name.
    """
    prefix = config.model_table_prefix
    dataset_name = Path(dataset_path).name.lower().replace("-", "_")
    return f"{prefix}_{dataset_name}_{model_name}"


def _ensure_table(
    config: BenchmarkConfig,
    model_name: str,
    vector_dim: int,
    dataset_path: str,
) -> lancedb.table.Table:
    """Ensure the LanceDB table exists with correct schema.

    Args:
        config: Benchmark configuration.
        model_name: Model name for table naming.
        vector_dim: Feature vector dimension.
        dataset_path: Dataset identifier or path.

    Returns:
        LanceDB table instance.
    """
    import pyarrow as pa
    from database import db_manager

    table_name = _get_table_name(config, model_name, dataset_path)

    # Build expected schema
    schema = pa.schema(
        [
            pa.field("id", pa.int32()),
            pa.field("label", pa.int32()),
            pa.field("vector", pa.list_(pa.float32(), vector_dim)),
        ]
    )

    db = db_manager.db
    existing_tables = db.list_tables().tables

    # Check if table exists and has correct schema
    if table_name in existing_tables:
        table = db.open_table(table_name)
        if table.schema != schema:
            console.print(
                f"[yellow]Table '{table_name}' schema mismatch, rebuilding.[/yellow]"
            )
            db.drop_table(table_name)
            table = db.create_table(table_name, schema=schema)
    else:
        table = db.create_table(table_name, schema=schema)

    return table


def _print_benchmark_info(
    config: BenchmarkConfig,
    vector_dim: int,
    table_name: str,
    table_count: int,
    dataset_path: str = "",
) -> None:
    """Print benchmark configuration info using Rich table.

    Args:
        config: Benchmark configuration.
        vector_dim: Feature vector dimension.
        table_name: Database table name.
        table_count: Number of entries in the table.
        dataset_path: Current dataset identifier or path.
    """
    table = Table(title="Benchmark Configuration", show_header=False)
    table.add_column("Key", style="cyan", no_wrap=True)
    table.add_column("Value", style="magenta")

    table.add_row("Dataset", dataset_path)
    table.add_row("Model Output Dimension", str(vector_dim))
    table.add_row("Table Name", table_name)
    table.add_row("Table Entries", str(table_count))

    console.print(table)


def _print_benchmark_results(results: dict[str, Any]) -> None:
    """Print benchmark results using Rich table.

    Args:
        results: Final benchmark metrics for a single dataset.
    """
    table = Table(title="Benchmark Results", show_header=False)
    table.add_column("Metric", style="cyan", no_wrap=True)
    table.add_column("Value", style="green")

    # Print recalls first
    recalls = results.get("recalls", {})
    for k_str, v in recalls.items():
        if isinstance(v, float):
            table.add_row(k_str, f"{v:.4f}")

    for key, value in results.items():
        if key == "recalls":
            continue
        if isinstance(value, (list, dict)):
            continue
        if isinstance(value, float):
            table.add_row(key, f"{value:.4f}")
            continue
        table.add_row(key, str(value))

    console.print(table)


def _save_benchmark_outputs(
    all_results: dict[str, dict[str, Any]],
    output_root: Path,
    model_name: str,
) -> Path:
    """Save multi-dataset benchmark results to disk.

    Writes to ``output_root/{run_id}/``:
        - summary.md — overall summary with per-dataset recall table
        - metrics.csv — one row per (dataset, K) combination
        - {dataset_name}/predictions.csv — per-sample predictions
        - {dataset_name}/confusion_matrix.csv — row-normalized (Top-1)

    Args:
        all_results: Dict mapping dataset name to evaluate() result dict.
        output_root: Parent directory (e.g. ``outputs/benchmark``).
        model_name: Model identifier for the run_id.

    Returns:
        Path to the created run directory.
    """
    run_id = f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-{model_name}"
    out_dir = output_root / run_id
    out_dir.mkdir(parents=True, exist_ok=True)

    # --- metrics.csv: one row per (dataset, k) ---
    csv_keys = ["model", "run_id", "dataset", "k", "recall@k", "total", "num_classes"]
    all_csv_rows: list[dict[str, Any]] = []

    # --- summary.md lines ---
    md_lines = [
        "# Benchmark Results",
        "",
        f"- **run_id**: `{run_id}`",
        f"- **model**: `{model_name}`",
        "",
        "## Per-Dataset Recall",
        "",
    ]

    for dataset_name, results in all_results.items():
        recalls = results.get("recalls", {})
        total = results.get("total", 0)
        num_classes = results.get("num_classes", 0)

        md_lines.append(f"### {dataset_name}")
        md_lines.append("")
        md_lines.append(f"- total queries: `{total}`")
        md_lines.append(f"- num_classes: `{num_classes}`")
        md_lines.append("")
        md_lines.append("| K | Recall@K |")
        md_lines.append("|---:|---:|")

        for k_str, recall_val in sorted(recalls.items()):
            k = int(k_str.replace("recall@", ""))
            md_lines.append(f"| {k} | {recall_val:.4f} |")
            all_csv_rows.append({
                "model": model_name,
                "run_id": run_id,
                "dataset": dataset_name,
                "k": k,
                "recall@k": recall_val,
                "total": total,
                "num_classes": num_classes,
            })

        md_lines.append("")

        # Per-dataset subdirectory
        dataset_dir = out_dir / dataset_name
        dataset_dir.mkdir(parents=True, exist_ok=True)

        # --- predictions.csv (per-sample) ---
        query_labels = results.get("query_labels", [])
        topk_labels = results.get("topk_labels", [])
        if query_labels and topk_labels:
            max_k = max(len(preds) for preds in topk_labels) if topk_labels else 1
            fieldnames = ["idx", "true_label"] + [
                f"top{i+1}_label" for i in range(max_k)
            ]
            with (dataset_dir / "predictions.csv").open(
                "w", newline="", encoding="utf-8"
            ) as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames)
                writer.writeheader()
                for idx, (true_label, preds) in enumerate(
                    zip(query_labels, topk_labels)
                ):
                    row_data: dict[str, Any] = {
                        "idx": idx,
                        "true_label": true_label,
                    }
                    for i in range(max_k):
                        row_data[f"top{i+1}_label"] = (
                            preds[i] if i < len(preds) else ""
                        )
                    writer.writerow(row_data)

        # --- confusion_matrix.csv ---
        cm = results.get("confusion_matrix")
        if cm is not None:
            with (dataset_dir / "confusion_matrix.csv").open(
                "w", newline="", encoding="utf-8"
            ) as fh:
                writer = csv.writer(fh)
                for row in cm:
                    writer.writerow([f"{v:.6f}" for v in row])

    # Write metrics.csv
    if all_csv_rows:
        with (out_dir / "metrics.csv").open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=csv_keys)
            writer.writeheader()
            for row in all_csv_rows:
                writer.writerow(row)

    # Write summary.md
    (out_dir / "summary.md").write_text("\n".join(md_lines) + "\n", encoding="utf-8")

    console.print(f"[green]Results saved to:[/green] {out_dir}")
    return out_dir


def run_benchmark(
    model: Any,
    processor: Any,
    config: BenchmarkConfig,
    model_config: ModelConfig | None = None,
    model_name: str = "model",
    output_root: Path | None = None,
) -> dict[str, Any]:
    """Run benchmark evaluation across one or more datasets.

    Workflow:
        1. Load model / processor once
        2. For each dataset in config.datasets:
            a. Create dataset, get splits
            b. Build evaluation database
            c. Evaluate with all K values in top_k_list
        3. Print per-dataset results
        4. Save all results to disk (if ``output_root`` provided)

    Args:
        model: Feature extraction model.
        processor: Image preprocessor.
        config: Benchmark configuration.
        model_config: Optional model configuration for task-owned loading.
        model_name: Model name for table naming.
        output_root: Optional root directory for saving results.

    Returns:
        Dict mapping dataset name to evaluation result dict.

    Raises:
        ValueError: If benchmark is not enabled in config.
    """
    # ── Resolve model / processor once ──────────────────────────
    task = _create_task(config, model_config)

    resolver = getattr(task, "prepare_benchmark", None)
    if callable(resolver):
        prepare_benchmark = cast(
            Callable[[Any, Any, str], tuple[Any, Any, str]],
            resolver,
        )
        model, processor, model_name = prepare_benchmark(model, processor, model_name)

    if model is None or processor is None:
        raise ValueError("Benchmark task did not provide a valid model and processor")

    # ── Infer vector dimension from first dataset ─────────────────
    first_cfg = config.datasets[0]
    first_dataset = create_dataset(first_cfg)
    first_train = first_dataset.get_train_split()
    sample = first_train[0]
    sample_image = sample["img"]

    from utils.feature_extractor import infer_vector_dim

    vector_dim = infer_vector_dim(processor, model, sample_image)
    console.print(f"[cyan]Model output dimension:[/cyan] {vector_dim}")

    # ── Evaluate each dataset ────────────────────────────────────
    all_results: dict[str, dict[str, Any]] = {}

    for dataset_cfg in config.datasets:
        console.print(
            f"\n[bold cyan]Dataset:[/bold cyan] "
            f"{dataset_cfg.source_type} - {dataset_cfg.path}"
        )

        dataset = create_dataset(dataset_cfg)
        train_dataset = dataset.get_train_split()
        test_dataset = dataset.get_test_split()

        if train_dataset is None or test_dataset is None:
            raise ValueError(
                f"Dataset {dataset_cfg.path} does not have train/test splits"
            )

        dataset_name = Path(dataset_cfg.path).name.lower()
        table = _ensure_table(config, model_name, vector_dim, dataset_cfg.path)
        table_name = _get_table_name(config, model_name, dataset_cfg.path)

        # Build database (skip if cached)
        table_count = table.count_rows()
        expected_count = len(train_dataset)

        if table_count == expected_count:
            console.print(
                f"[yellow]Table '{table_name}' already has {table_count} entries, "
                f"skipping database build.[/yellow]"
            )
        else:
            if table_count > 0:
                console.print(
                    f"[yellow]Table '{table_name}' has {table_count} entries "
                    f"(expected {expected_count}), rebuilding.[/yellow]"
                )
                # Rebuild: drop and recreate
                from database import db_manager
                db_manager.db.drop_table(table_name)
                table = _ensure_table(config, model_name, vector_dim, dataset_cfg.path)

            console.print(
                f"[cyan]Building database[/cyan] with {expected_count} training samples..."
            )
            task.build_database(
                model, processor, train_dataset, table, config.batch_size,
                label_column=dataset_cfg.label_column,
            )
            table_count = table.count_rows()

        _print_benchmark_info(
            config, vector_dim, table_name, table_count,
            dataset_path=dataset_cfg.path,
        )

        # Evaluate
        console.print(
            f"[cyan]Evaluating[/cyan] on {len(test_dataset)} test samples..."
        )
        results = task.evaluate(
            model, processor, test_dataset, table, config.batch_size,
            label_column=dataset_cfg.label_column,
        )
        all_results[dataset_name] = results

        # Print per-dataset results
        _print_benchmark_results(results)

    # ── Save outputs ─────────────────────────────────────────────
    if output_root is not None:
        _save_benchmark_outputs(all_results, output_root, model_name)

    return all_results