"""SeekDB backend implementation."""

from pathlib import Path
from typing import Any

import pyseekdb
from pyseekdb import Configuration, HNSWConfiguration, FulltextIndexConfig

from .base import BaseBackend, SearchHit, SearchResult
from . import register_backend


@register_backend('seekdb')
class SeekDBBackend(BaseBackend):
    """SeekDB (OceanBase) backend implementation."""

    def __init__(
        self,
        path: str | None = None,
        host: str | None = None,
        port: int = 2881,
        database: str = 'mykb',
        fulltext_analyzer: str = 'ik',  # ik | space | ngram
        **kwargs,
    ):
        self.database = database
        self.fulltext_analyzer = fulltext_analyzer
        self._path = str(Path(path).expanduser()) if path else None

        if path:
            # Embedded mode - need to create database first
            admin = pyseekdb.AdminClient(path=self._path)
            try:
                admin.create_database(database)
            except Exception:
                pass  # Database may already exist
            self.client = pyseekdb.Client(path=self._path, database=database)
        else:
            # Server mode
            self.client = pyseekdb.Client(host=host, port=port, database=database, **kwargs)

    def initialize(self, collection: str, dimension: int) -> None:
        """Create collection if not exists."""
        if self.client.has_collection(collection):
            return

        config = Configuration(
            hnsw=HNSWConfiguration(dimension=dimension, distance='cosine'),
            fulltext_config=FulltextIndexConfig(analyzer=self.fulltext_analyzer),
        )
        self.client.create_collection(
            name=collection,
            configuration=config,
            embedding_function=None,  # We manage embeddings ourselves
        )

    def add_documents(self, collection: str, documents: list[dict]) -> int:
        """Add or update documents."""
        if not documents:
            return 0

        coll = self.client.get_collection(collection, embedding_function=None)

        ids = []
        embeddings = []
        docs = []
        metadatas = []

        for d in documents:
            ids.append(d['id'])
            docs.append(d.get('content', ''))

            # Extract embedding (if present)
            vectors = d.get('_vectors', {})
            emb = vectors.get('gemma') if vectors else None
            embeddings.append(emb)

            # Rest of fields as metadata
            meta = {k: v for k, v in d.items() if k not in ('id', 'content', '_vectors')}
            metadatas.append(meta)

        # Check if any embeddings are present
        has_embeddings = any(e is not None for e in embeddings)

        coll.upsert(ids=ids, documents=docs, embeddings=embeddings if has_embeddings else None, metadatas=metadatas)
        return len(ids)

    def delete_documents(self, collection: str, ids: list[str]) -> int:
        """Delete documents by ID."""
        if not ids:
            return 0
        coll = self.client.get_collection(collection, embedding_function=None)
        coll.delete(ids=ids)
        return len(ids)

    def update_embeddings(self, collection: str, id_vectors: dict[str, list[float]]) -> int:
        """Update document embeddings."""
        if not id_vectors:
            return 0
        coll = self.client.get_collection(collection, embedding_function=None)
        ids = list(id_vectors.keys())
        embeddings = list(id_vectors.values())
        coll.update(ids=ids, embeddings=embeddings)
        return len(ids)

    def bm25_search(
        self,
        collection: str,
        query: str,
        limit: int = 10,
        filter: dict | None = None,
    ) -> SearchResult:
        """Pure BM25 full-text search via get() + where_document."""
        coll = self.client.get_collection(collection, embedding_function=None)

        results = coll.get(
            where_document={'$contains': query}, where=filter, limit=limit, include=['documents', 'metadatas']
        )

        return self._convert_get_result(results)

    def vector_search(
        self,
        collection: str,
        vector: list[float],
        limit: int = 10,
        filter: dict | None = None,
    ) -> SearchResult:
        """Pure vector search via query()."""
        coll = self.client.get_collection(collection, embedding_function=None)

        results = coll.query(
            query_embeddings=[vector], where=filter, n_results=limit, include=['documents', 'metadatas']
        )

        return self._convert_query_result(results)

    def hybrid_search(
        self,
        collection: str,
        query: str,
        vector: list[float],
        limit: int = 10,
        semantic_ratio: float = 0.5,
        filter: dict | None = None,
    ) -> SearchResult:
        """Hybrid search using RRF fusion."""
        coll = self.client.get_collection(collection, embedding_function=None)

        # Build hybrid search params
        query_param: dict[str, Any] = {'where_document': {'$contains': query}}
        if filter:
            query_param['where'] = filter

        knn_param: dict[str, Any] = {
            'query_embeddings': [vector],
            'n_results': limit * 2,  # Get more candidates for RRF
        }
        if filter:
            knn_param['where'] = filter

        results = coll.hybrid_search(
            query=query_param,
            knn=knn_param,
            rank={'rrf': {'rank_constant': 60}},
            n_results=limit,
            include=['documents', 'metadatas'],
        )

        return self._convert_hybrid_result(results)

    def get_document(self, collection: str, id: str) -> dict | None:
        """Get a single document by ID."""
        coll = self.client.get_collection(collection, embedding_function=None)
        results = coll.get(ids=[id], include=['documents', 'metadatas', 'embeddings'])

        if not results.get('ids'):
            return None

        return {
            'id': results['ids'][0],
            'content': results['documents'][0] if results.get('documents') else None,
            'metadata': results['metadatas'][0] if results.get('metadatas') else {},
            'embedding': results['embeddings'][0] if results.get('embeddings') else None,
        }

    def get_documents(
        self,
        collection: str,
        limit: int = 1000,
        offset: int = 0,
        filter: str | None = None,
    ) -> list[dict]:
        """Get documents with pagination."""
        coll = self.client.get_collection(collection, embedding_function=None)

        # SeekDB get() uses where filter instead of string filter
        results = coll.get(limit=limit, offset=offset, include=['documents', 'metadatas'])

        docs = []
        ids = results.get('ids', [])
        documents = results.get('documents', [])
        metadatas = results.get('metadatas', [])

        for i, id_ in enumerate(ids):
            doc = {
                'id': id_,
                'content': documents[i] if documents else None,
            }
            if metadatas and metadatas[i]:
                doc.update(metadatas[i])
            docs.append(doc)

        return docs

    def get_stats(self, collection: str) -> dict:
        """Get collection statistics."""
        coll = self.client.get_collection(collection, embedding_function=None)
        return {
            'total_documents': coll.count(),
            'dimension': getattr(coll, 'dimension', None),
            'distance': getattr(coll, 'distance', None),
        }

    def health_check(self) -> dict:
        """Health check."""
        try:
            collections = self.client.list_collections()
            return {
                'status': 'ok',
                'backend': 'seekdb',
                'collections': len(collections),
            }
        except Exception as e:
            return {'status': 'error', 'message': str(e)}

    def _convert_get_result(self, results: dict) -> SearchResult:
        """Convert get() return format to SearchResult."""
        hits = []
        for i, id_ in enumerate(results.get('ids', [])):
            doc = {}
            if results.get('documents'):
                doc['content'] = results['documents'][i]
            if results.get('metadatas') and results['metadatas'][i]:
                doc.update(results['metadatas'][i])

            hits.append(
                SearchHit(
                    id=id_,
                    score=1.0,  # get() has no score
                    document=doc,
                )
            )
        return SearchResult(hits=hits, processing_time_ms=0)

    def _convert_query_result(self, results: dict) -> SearchResult:
        """Convert query() return format to SearchResult."""
        hits = []
        # query returns List[List[...]], take first query's results
        ids = results.get('ids', [[]])[0]
        distances = results.get('distances', [[]])[0]
        documents = results.get('documents', [[]])[0] if results.get('documents') else [None] * len(ids)
        metadatas = results.get('metadatas', [[]])[0] if results.get('metadatas') else [{}] * len(ids)

        for i, id_ in enumerate(ids):
            doc = {'content': documents[i]}
            if metadatas[i]:
                doc.update(metadatas[i])

            # Convert distance to similarity score
            score = 1.0 / (1.0 + float(distances[i])) if distances else 1.0
            hits.append(SearchHit(id=id_, score=score, document=doc))

        return SearchResult(hits=hits, processing_time_ms=0)

    def _convert_hybrid_result(self, results: dict) -> SearchResult:
        """Convert hybrid_search() return format to SearchResult."""
        # hybrid_search returns nested lists like query()
        hits = []
        ids = results.get('ids', [[]])[0]
        distances = results.get('distances', [[]])[0] if results.get('distances') else [0.0] * len(ids)
        documents = results.get('documents', [[]])[0] if results.get('documents') else [None] * len(ids)
        metadatas = results.get('metadatas', [[]])[0] if results.get('metadatas') else [{}] * len(ids)

        for i, id_ in enumerate(ids):
            doc = {'content': documents[i]}
            if metadatas[i]:
                doc.update(metadatas[i])

            # Convert Decimal to float if needed
            dist = float(distances[i]) if distances[i] is not None else 0.0
            score = 1.0 / (1.0 + dist)
            hits.append(SearchHit(id=id_, score=score, document=doc))

        return SearchResult(hits=hits, processing_time_ms=0)
