# src/index/faiss_index.py
import faiss
import numpy as np
import pickle
import os
from typing import List, Tuple, Optional, Dict
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)


@dataclass
class IndexConfig:
    dimension: int = 384
    index_type: str = "hnsw"  # flat, ivf, hnsw
    hnsw_m: int = 32
    hnsw_ef_construction: int = 200
    hnsw_ef_search: int = 100
    nprobe: int = 10  # for IVF


class FAISSIndexBuilder:
    def __init__(self, config: IndexConfig = None):
        self.config = config or IndexConfig()
        self.index = None
        self.id_map = {}  # internal_id -> external_id
        self.reverse_map = {}  # external_id -> internal_id
        self.next_id = 0

    def build_index(self, embeddings: np.ndarray, ids: List[str]):
        """Build index from scratch with pre-computed embeddings."""
        n_vectors = len(embeddings)
        dim = embeddings.shape[1]

        if self.config.index_type == "flat":
            self.index = faiss.IndexFlatIP(dim)  # Inner product for cosine similarity
        elif self.config.index_type == "hnsw":
            self.index = faiss.IndexHNSWFlat(dim, self.config.hnsw_m)
            self.index.hnsw.efConstruction = self.config.hnsw_ef_construction
            self.index.hnsw.efSearch = self.config.hnsw_ef_search
        elif self.config.index_type == "ivf":
            quantizer = faiss.IndexFlatIP(dim)
            self.index = faiss.IndexIVFFlat(quantizer, dim, n_vectors // 10)
            self.index.nprobe = self.config.nprobe
            self.index.train(embeddings.astype(np.float32))

        # Normalize for cosine similarity
        faiss.normalize_L2(embeddings.astype(np.float32))

        self.index.add(embeddings.astype(np.float32))

        # Build ID maps
        for i, ext_id in enumerate(ids):
            self.id_map[i] = ext_id
            self.reverse_map[ext_id] = i
        self.next_id = len(ids)

        logger.info(f"Built {self.config.index_type} index with {n_vectors} vectors, dim={dim}")

    def search(
        self, query: np.ndarray, top_k: int = 10,
        filter_ids: Optional[List[str]] = None
    ) -> List[Tuple[str, float]]:
        """Search for nearest neighbors."""
        query = query.reshape(1, -1).astype(np.float32)
        faiss.normalize_L2(query)

        # Search more if filtering
        search_k = top_k * 3 if filter_ids else top_k
        scores, indices = self.index.search(query, min(search_k, self.index.ntotal))

        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx == -1:
                continue
            ext_id = self.id_map.get(idx)
            if ext_id is None:
                continue
            if filter_ids and ext_id not in filter_ids:
                continue
            results.append((ext_id, float(score)))
            if len(results) >= top_k:
                break

        return results

    def add_vectors(self, embeddings: np.ndarray, ids: List[str]):
        """Add new vectors to the index."""
        faiss.normalize_L2(embeddings.astype(np.float32))
        self.index.add(embeddings.astype(np.float32))
        for ext_id in ids:
            self.id_map[self.next_id] = ext_id
            self.reverse_map[ext_id] = self.next_id
            self.next_id += 1

    def save(self, path: str):
        """Save index and mappings to disk."""
        os.makedirs(path, exist_ok=True)
        faiss.write_index(self.index, os.path.join(path, "index.faiss"))
        with open(os.path.join(path, "mappings.pkl"), "wb") as f:
            pickle.dump({
                "id_map": self.id_map,
                "reverse_map": self.reverse_map,
                "next_id": self.next_id,
                "config": self.config,
            }, f)
        logger.info(f"Saved index to {path}")

    def load(self, path: str):
        """Load index and mappings from disk."""
        self.index = faiss.read_index(os.path.join(path, "index.faiss"))
        with open(os.path.join(path, "mappings.pkl"), "rb") as f:
            data = pickle.load(f)
            self.id_map = data["id_map"]
            self.reverse_map = data["reverse_map"]
            self.next_id = data["next_id"]
            self.config = data["config"]
        logger.info(f"Loaded index from {path} with {self.index.ntotal} vectors")

# src/embeddings/embedding_service.py
import numpy as np
from typing import List, Union
from sentence_transformers import SentenceTransformer
import openai
import hashlib
import json


class EmbeddingService:
    def __init__(self, model_type: str = "sentence-transformer", model_name: str = None):
        self.model_type = model_type
        if model_type == "sentence-transformer":
            self.model = SentenceTransformer(model_name or "all-MiniLM-L6-v2")
            self.dimension = self.model.get_sentence_embedding_dimension()
        elif model_type == "openai":
            self.model = None
            self.dimension = 1536

    def embed_texts(self, texts: List[str], batch_size: int = 256) -> np.ndarray:
        if self.model_type == "sentence-transformer":
            return self.model.encode(texts, batch_size=batch_size, show_progress_bar=True)
        elif self.model_type == "openai":
            embeddings = []
            for i in range(0, len(texts), 2048):
                batch = texts[i:i+2048]
                response = openai.embeddings.create(input=batch, model="text-embedding-3-small")
                embeddings.extend([e.embedding for e in response.data])
            return np.array(embeddings)

    def embed_query(self, query: str) -> np.ndarray:
        return self.embed_texts([query])[0]

    def embed_image(self, image_path: str) -> np.ndarray:
        # Placeholder for CLIP-based image embedding
        # In production, use open_clip or CLIP model
        raise NotImplementedError("Image embedding not yet implemented")

# src/search/hybrid_search.py
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
import numpy as np
from rank_bm25 import BM25Okapi
import jieba


@dataclass
class SearchResult:
    id: str
    score: float
    content: str
    metadata: Dict
    search_type: str  # semantic, keyword, hybrid


class HybridSearchEngine:
    def __init__(self, alpha: float = 0.7):
        """alpha: weight for semantic search (1-alpha for keyword)."""
        self.alpha = alpha
        self.vector_index = None
        self.bm25 = None
        self.documents = {}
        self.corpus_tokens = []

    def index_documents(self, documents: List[Dict]):
        """Index documents for both semantic and keyword search."""
        self.documents = {doc["id"]: doc for doc in documents}

        # Build BM25 index
        self.corpus_tokens = [
            list(jieba.cut(doc["content"])) for doc in documents
        ]
        self.bm25 = BM25Okapi(self.corpus_tokens)

    def search(
        self, query: str, query_embedding: np.ndarray,
        top_k: int = 10, filter_fn=None
    ) -> List[SearchResult]:
        """Hybrid search combining semantic and keyword results."""

        # Semantic search
        semantic_results = self.vector_index.search(query_embedding, top_k=top_k * 2)

        # Keyword search
        query_tokens = list(jieba.cut(query))
        bm25_scores = self.bm25.get_scores(query_tokens)
        bm25_top_idx = np.argsort(bm25_scores)[::-1][:top_k * 2]
        keyword_results = [
            (list(self.documents.keys())[idx], bm25_scores[idx]) for idx in bm25_top_idx
        ]

        # Merge results with Reciprocal Rank Fusion
        merged = {}
        k = 60  # RRF constant

        for rank, (doc_id, score) in enumerate(semantic_results):
            merged[doc_id] = merged.get(doc_id, 0) + self.alpha / (k + rank + 1)

        for rank, (doc_id, score) in enumerate(keyword_results):
            merged[doc_id] = merged.get(doc_id, 0) + (1 - self.alpha) / (k + rank + 1)

        # Sort and return top_k
        sorted_results = sorted(merged.items(), key=lambda x: x[1], reverse=True)[:top_k]

        return [
            SearchResult(
                id=doc_id, score=score,
                content=self.documents[doc_id]["content"],
                metadata=self.documents[doc_id].get("metadata", {}),
                search_type="hybrid",
            )
            for doc_id, score in sorted_results
        ]

# src/api/main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import List, Optional, Dict
from contextlib import asynccontextmanager
import numpy as np
import time


@asynccontextmanager
async def lifespan(app: FastAPI):
    from src.index.faiss_index import FAISSIndexBuilder
    from src.embeddings.embedding_service import EmbeddingService
    from src.search.hybrid_search import HybridSearchEngine

    app.state.embedder = EmbeddingService("sentence-transformer")
    app.state.search_engine = HybridSearchEngine(alpha=0.7)
    # Load pre-built index...
    yield


app = FastAPI(title="Vector Search API", version="1.0.0", lifespan=lifespan)


class SearchRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=1000)
    top_k: int = Field(10, ge=1, le=100)
    filters: Optional[Dict] = None
    search_mode: str = Field("hybrid", regex="^(semantic|keyword|hybrid)$")


class SearchResponse(BaseModel):
    results: List[Dict]
    total_results: int
    latency_ms: float


@app.post("/api/v1/search", response_model=SearchResponse)
async def search(request: SearchRequest):
    start = time.time()

    # Generate query embedding
    query_emb = app.state.embedder.embed_query(request.query)

    # Search
    results = app.state.search_engine.search(
        query=request.query,
        query_embedding=query_emb,
        top_k=request.top_k,
    )

    latency = (time.time() - start) * 1000

    return SearchResponse(
        results=[{"id": r.id, "score": r.score, "content": r.content[:200]} for r in results],
        total_results=len(results),
        latency_ms=latency,
    )


@app.post("/api/v1/index")
async def index_documents(documents: List[Dict]):
    app.state.search_engine.index_documents(documents)
    return {"indexed": len(documents)}