"""Retrieval task for benchmark evaluation (Recall@K)."""

from typing import TYPE_CHECKING, Any, cast

import lancedb
import pyarrow as pa
import torch
import torch.nn.functional as F
from benchmarks.base import BaseBenchmarkTask
from benchmarks.tasks.registry import RegisterTask
from compressors.model_loader import get_dino_dim, load_dino_model, load_hash_compressor
from configs import cfg_manager
from rich.progress import track
from torch import nn
from torch.utils.data import DataLoader
from transformers import BitImageProcessor
from utils.feature_extractor import extract_batch_features, infer_vector_dim

if TYPE_CHECKING:
    from compressors.hash_compressor import HashCompressor


class RetrievalEncoder(nn.Module):
    """Benchmark encoder for DINO and optional hash compression."""

    def __init__(
        self,
        dino: nn.Module,
        compressor: "HashCompressor | None" = None,
    ) -> None:
        """Initialize retrieval encoder.

        Args:
            dino: DINO backbone used for feature extraction.
            compressor: Optional hash compressor for recall evaluation.
        """
        super().__init__()
        self.dino: nn.Module = dino
        self.compressor: HashCompressor | None = compressor

    def forward(self, inputs: Any) -> torch.Tensor:
        """Encode processor inputs into benchmark vectors.

        Args:
            inputs: Batched processor outputs.

        Returns:
            Float tensor used for LanceDB insertion and retrieval.
        """
        outputs = self.dino(**inputs)
        tokens = outputs.last_hidden_state

        if self.compressor is None:
            features = tokens.mean(dim=1)
            return F.normalize(features, dim=-1)

        bits = self.compressor.encode(tokens)
        return bits.to(dtype=torch.float32)


def _build_eval_schema(vector_dim: int) -> pa.Schema:
    """Build PyArrow schema for evaluation database table.

    Args:
        vector_dim: Feature vector dimension.

    Returns:
        PyArrow schema with id, label, and vector fields.
    """
    return pa.schema(
        [
            pa.field("id", pa.int32()),
            pa.field("label", pa.int32()),
            pa.field("vector", pa.list_(pa.float32(), vector_dim)),
        ]
    )


def _establish_eval_database(
    processor: BitImageProcessor,
    model: nn.Module,
    table: lancedb.table.Table,
    dataloader: DataLoader[Any],
) -> None:
    """Extract features from training images and store them in a database table.

    Args:
        processor: Image preprocessor.
        model: Feature extraction model.
        table: LanceDB table to store features.
        dataloader: DataLoader for the training dataset.
    """
    # Extract all features using the utility function
    all_features = extract_batch_features(processor, model, dataloader)
    config = cfg_manager.get()

    # Store features to database
    global_idx = 0
    for batch in track(dataloader, description="Storing eval database"):
        labels = batch[config.benchmark.dataset.label_column]
        labels_list = labels.tolist()
        batch_size = len(labels_list)

        table.add(
            [
                {
                    "id": global_idx + j,
                    "label": labels_list[j],
                    "vector": all_features[global_idx + j].detach().cpu().numpy(),
                }
                for j in range(batch_size)
            ]
        )
        global_idx += batch_size


def _evaluate_recall(
    processor: BitImageProcessor,
    model: nn.Module,
    table: lancedb.table.Table,
    dataloader: DataLoader[Any],
    top_k: int,
) -> tuple[int, int]:
    """Evaluate Recall@K by searching the database for each test image.

    Args:
        processor: Image preprocessor.
        model: Feature extraction model.
        table: LanceDB table to search against.
        dataloader: DataLoader for the test dataset.
        top_k: Number of top results to retrieve.

    Returns:
        A tuple of (correct_count, total_count).
    """
    # Extract all features using the utility function
    all_features = extract_batch_features(processor, model, dataloader)
    config = cfg_manager.get()

    correct = 0
    total = 0
    feature_idx = 0

    for batch in track(dataloader, description=f"Evaluating Recall@{top_k}"):
        labels = batch[config.benchmark.dataset.label_column]
        labels_list = labels.tolist()

        for j in range(len(labels_list)):
            feature = all_features[feature_idx + j].tolist()
            true_label = labels_list[j]

            results = (
                table.search(feature)
                .select(["label", "_distance"])
                .limit(top_k)
                .to_polars()
            )

            retrieved_labels = results["label"].to_list()
            if true_label in retrieved_labels:
                correct += 1
            total += 1

        feature_idx += len(labels_list)

    return correct, total


@RegisterTask("retrieval")
class RetrievalTask(BaseBenchmarkTask):
    """Retrieval evaluation task (Recall@K)."""

    def __init__(
        self,
        top_k: int = 10,
        dino_model: str = "facebook/dinov2-large",
        compression_dim: int = 512,
        compressor_path: str | None = None,
    ):
        """Initialize retrieval task.

        Args:
            top_k: Number of top results to retrieve for recall calculation.
            dino_model: DINO model name used for feature extraction.
            compression_dim: Output dimension of the hash compressor.
            compressor_path: Optional path to trained hash compressor weights.
        """
        super().__init__(top_k=top_k)
        self.top_k = top_k
        self.dino_model = dino_model
        self.compression_dim = compression_dim
        self.compressor_path = compressor_path
        self._processor: BitImageProcessor | None = None
        self._model: nn.Module | None = None
        self._model_name = "hash_compressor" if compressor_path else "dinov2"

    def prepare_benchmark(
        self,
        model: Any,
        processor: Any,
        model_name: str = "model",
    ) -> tuple[nn.Module, BitImageProcessor, str]:
        """Resolve benchmark resources for this task.

        Args:
            model: Optional pre-built model from the caller.
            processor: Optional pre-built processor from the caller.
            model_name: Fallback table model name.

        Returns:
            Tuple of benchmark model, processor, and resolved model name.
        """
        if model is not None and processor is not None:
            return (
                cast(nn.Module, model),
                cast(BitImageProcessor, processor),
                model_name,
            )

        self._ensure_resources_loaded()
        return (
            cast(nn.Module, self._model),
            cast(BitImageProcessor, self._processor),
            self._model_name,
        )

    def _ensure_resources_loaded(self) -> None:
        """Lazy-load retrieval benchmark resources."""
        if self._processor is not None and self._model is not None:
            return

        processor, dino = load_dino_model(self.dino_model)

        compressor = None
        if self.compressor_path is not None:
            compressor = load_hash_compressor(
                input_dim=get_dino_dim(self.dino_model),
                hash_bits=self.compression_dim,
                compressor_path=self.compressor_path,
            )
            compressor.eval()

        self._processor = processor
        self._model = RetrievalEncoder(dino=dino, compressor=compressor)
        self._model.eval()

    def build_database(
        self,
        model: Any,
        processor: Any,
        train_dataset: Any,
        table: lancedb.table.Table,
        batch_size: int,
    ) -> None:
        """Build the evaluation database from training data.

        Args:
            model: Feature extraction model.
            processor: Image preprocessor.
            train_dataset: Training dataset.
            table: LanceDB table to store features.
            batch_size: Batch size for DataLoader.
        """
        # Get a sample image to infer vector dimension
        sample = train_dataset[0]
        sample_image = sample["img"]

        vector_dim = infer_vector_dim(processor, model, sample_image)
        expected_schema = _build_eval_schema(vector_dim)

        # Check schema compatibility
        if table.schema != expected_schema:
            raise ValueError(
                f"Table schema mismatch. Expected: {expected_schema}, "
                f"Got: {table.schema}"
            )

        # Build database
        train_loader = DataLoader(
            train_dataset.with_format("torch"),
            batch_size=batch_size,
            shuffle=False,
            num_workers=4,
        )
        _establish_eval_database(processor, model, table, train_loader)

    def evaluate(
        self,
        model: Any,
        processor: Any,
        test_dataset: Any,
        table: lancedb.table.Table,
        batch_size: int,
    ) -> dict[str, Any]:
        """Evaluate the model on the test dataset.

        Args:
            model: Feature extraction model.
            processor: Image preprocessor.
            test_dataset: Test dataset.
            table: LanceDB table to search against.
            batch_size: Batch size for DataLoader.

        Returns:
            Dictionary containing evaluation results with keys:
                - accuracy: Recall@K accuracy (0.0 ~ 1.0)
                - correct: Number of correct predictions
                - total: Total number of test samples
                - top_k: The K value used
        """
        test_loader = DataLoader(
            test_dataset.with_format("torch"),
            batch_size=batch_size,
            shuffle=False,
            num_workers=4,
        )
        correct, total = _evaluate_recall(
            processor, model, table, test_loader, self.top_k
        )

        accuracy = correct / total if total > 0 else 0.0

        return {
            "accuracy": accuracy,
            "correct": correct,
            "total": total,
            "top_k": self.top_k,
        }