import logging
from typing import List, Optional, Tuple

import numpy as np
import faiss

from backend.app.config import SearchConfig

logger = logging.getLogger(__name__)


class DualFAISSIndex:
    """
    Two parallel FAISS indices (image + text) fused via Reciprocal Rank Fusion.
    """

    def __init__(self, dim: int, config: SearchConfig):
        self.dim = dim
        self.config = config
        self.image_index = None
        self.text_index = None

    def _create_index(self, n_vectors: int) -> faiss.Index:
        if n_vectors < 5000:
            logger.info(f"Using IndexFlatIP (exact, n={n_vectors:,})")
            return faiss.IndexFlatIP(self.dim)

        n_clusters = min(self.config.n_clusters, max(16, n_vectors // 40))
        logger.info(f"Using IndexIVFFlat (n={n_vectors:,}, clusters={n_clusters})")
        quantizer = faiss.IndexFlatIP(self.dim)
        index = faiss.IndexIVFFlat(
            quantizer, self.dim, n_clusters, faiss.METRIC_INNER_PRODUCT
        )
        return index

    def build(self, image_embeddings: np.ndarray, text_embeddings: np.ndarray):
        image_embeddings = image_embeddings.astype(np.float32)
        text_embeddings = text_embeddings.astype(np.float32)

        assert image_embeddings.shape == text_embeddings.shape, (
            f"Shape mismatch: images {image_embeddings.shape} vs text {text_embeddings.shape}"
        )

        n = image_embeddings.shape[0]

        logger.info("Building image FAISS index...")
        self.image_index = self._create_index(n)
        if hasattr(self.image_index, 'train'):
            try:
                self.image_index.train(image_embeddings)
            except Exception:
                self.image_index = faiss.IndexFlatIP(self.dim)
        self.image_index.add(image_embeddings)

        logger.info("Building text FAISS index...")
        self.text_index = self._create_index(n)
        if hasattr(self.text_index, 'train'):
            try:
                self.text_index.train(text_embeddings)
            except Exception:
                self.text_index = faiss.IndexFlatIP(self.dim)
        self.text_index.add(text_embeddings)

        logger.info(
            f"Dual index built: {self.image_index.ntotal:,} image, "
            f"{self.text_index.ntotal:,} text vectors"
        )

    def search_image_index(self, query: np.ndarray, top_k: int):
        q = query.astype(np.float32).reshape(1, -1)
        if hasattr(self.image_index, 'nprobe'):
            self.image_index.nprobe = self.config.n_probe
        return self.image_index.search(q, top_k)

    def search_text_index(self, query: np.ndarray, top_k: int):
        q = query.astype(np.float32).reshape(1, -1)
        if hasattr(self.text_index, 'nprobe'):
            self.text_index.nprobe = self.config.n_probe
        return self.text_index.search(q, top_k)

    def search_fused(
        self, query: np.ndarray, top_k: int,
        image_weight: Optional[float] = None,
        text_weight: Optional[float] = None,
    ) -> Tuple[List[int], List[float]]:
        iw = image_weight or self.config.image_index_weight
        tw = text_weight or self.config.text_index_weight
        rrf_k = self.config.rrf_k

        broad_k = min(top_k * 3, self.image_index.ntotal)

        _, img_ids = self.search_image_index(query, broad_k)
        _, txt_ids = self.search_text_index(query, broad_k)
        img_ids = img_ids[0]
        txt_ids = txt_ids[0]

        img_rank = {int(idx): rank + 1 for rank, idx in enumerate(img_ids) if idx >= 0}
        txt_rank = {int(idx): rank + 1 for rank, idx in enumerate(txt_ids) if idx >= 0}

        all_candidates = set(img_rank.keys()) | set(txt_rank.keys())
        scores = {}
        for idx in all_candidates:
            score = 0.0
            if idx in img_rank:
                score += iw / (rrf_k + img_rank[idx])
            if idx in txt_rank:
                score += tw / (rrf_k + txt_rank[idx])
            scores[idx] = score

        ranked = sorted(scores.items(), key=lambda x: -x[1])[:top_k]
        return [r[0] for r in ranked], [r[1] for r in ranked]

    def save(self, image_path: str, text_path: str):
        faiss.write_index(self.image_index, image_path)
        faiss.write_index(self.text_index, text_path)
        logger.info(f"Saved dual index to {image_path} and {text_path}")

    def load(self, image_path: str, text_path: str):
        self.image_index = faiss.read_index(image_path)
        self.text_index = faiss.read_index(text_path)
        logger.info(
            f"Loaded dual index: {self.image_index.ntotal:,} image, "
            f"{self.text_index.ntotal:,} text vectors"
        )


__all__ = ["DualFAISSIndex"]