"""Search functionality for docsctx (local and global)."""

from __future__ import annotations

import struct
from typing import Sequence


from docsctx.db.models import Document, DocumentChunk, DocumentEntity
from docsctx.db.repository import DocsRepository
from docsctx.embedder import BaseEmbedder, cosine_similarity
from docsctx.provider import BaseProvider

try:
    from docsctx.vectorstore import DocsVectorStore
except ImportError:
    DocsVectorStore = None  # type: ignore[misc,assignment]


def decode_vector(data: bytes) -> list[float]:
    """Unpack a bytes blob into a float list."""
    count = len(data) // 4
    return list(struct.unpack(f"<{count}f", data))


class LocalSearch:
    """Vector similarity search combined with graph walk (BFS).

    When a ``DocsVectorStore`` is provided, HNSW ANN search via ChromaDB is
    used (sub-ms at scale). Otherwise falls back to brute-force in the
    repository.
    """

    def __init__(
        self,
        repository: DocsRepository,
        embedder: BaseEmbedder,
        vectorstore: DocsVectorStore | None = None,
    ) -> None:
        self.repository = repository
        self.embedder = embedder
        self.vectorstore = vectorstore

    def search(
        self,
        query: str,
        top_k: int = 10,
        graph_depth: int = 2,
        min_score: float = 0.3,
        version: str | None = None,
    ) -> dict[str, object]:
        """
        Search for chunks via vector similarity, then walk graph from top entities.

        Args:
            query: Search query
            top_k: Number of top chunks to return
            graph_depth: Max depth for BFS graph walk from entities
            min_score: Minimum cosine similarity score
            version: Optional version filter (for version-aware search)

        Returns:
            Dict with chunks, entities, and relationships
        """
        # Embed the query
        query_embedding = self.embedder.embed_query(query)

        # Try ChromaDB HNSW first, then fall back to brute-force
        if self.vectorstore is not None:
            chroma_results = self.vectorstore.search_chunks(
                query_embedding, top_k=top_k, version=version,
            )
            chunks = [
                r for r in chroma_results if r["score"] >= min_score
            ]
        else:
            # Legacy brute-force path
            chunks = self.repository.search_chunks_by_embedding(
                query_embedding,
                limit=top_k,
                min_score=min_score,
            )

        # Collect entities from top chunks
        entity_ids = set()
        for chunk in chunks:
            doc_entities = self.repository.list_entities(chunk["document_id"])
            for entity in doc_entities:
                entity_ids.add(entity["id"])

        # Graph walk from entities (BFS up to graph_depth)
        graph_entities = self._graph_walk(entity_ids, depth=graph_depth)

        return {
            "query": query,
            "chunks": chunks,
            "entities": graph_entities,
            "count": len(chunks),
        }

    def _graph_walk(
        self,
        start_entity_ids: set[int],
        depth: int = 2,
    ) -> list[dict]:
        """BFS walk from entity set."""
        visited = set(start_entity_ids)
        frontier = list(start_entity_ids)
        depth_map = {eid: 0 for eid in frontier}

        for _ in range(depth):
            next_frontier = []
            for entity_id in frontier:
                # Get neighbors via relationships
                entity = self.repository.get_entity(entity_id)
                if not entity:
                    continue

                rels_from = self.repository.list_relationships_from_entity(entity["name"])
                rels_to = self.repository.list_relationships_to_entity(entity["name"])

                for rel in rels_from + rels_to:
                    target_name = rel["target"] if rel["source"] == entity["name"] else rel["source"]
                    target_entity = self.repository.get_entity_by_name(target_name)
                    if target_entity and target_entity["id"] not in visited:
                        visited.add(target_entity["id"])
                        depth_map[target_entity["id"]] = depth_map[entity_id] + 1
                        next_frontier.append(target_entity["id"])

            frontier = next_frontier
            if not frontier:
                break

        # Return all visited entities
        return [self.repository.get_entity(entity_id) for entity_id in visited if self.repository.get_entity(entity_id)]


class GlobalSearch:
    """Community-level search with synthesis.

    When a ``DocsVectorStore`` and ``embedder`` are provided, uses semantic
    HNSW search over community summaries. Otherwise falls back to keyword
    matching.
    """

    def __init__(
        self,
        repository: DocsRepository,
        provider: BaseProvider | None = None,
        vectorstore: DocsVectorStore | None = None,
        embedder: BaseEmbedder | None = None,
    ) -> None:
        self.repository = repository
        self.provider = provider
        self.vectorstore = vectorstore
        self.embedder = embedder

    def search(self, query: str, top_k: int = 5) -> dict[str, object]:
        """
        Search across community summaries.

        Args:
            query: Search query
            top_k: Number of top communities

        Returns:
            Dict with communities and entities
        """
        # Try semantic search via ChromaDB first
        if self.vectorstore is not None and self.embedder is not None:
            query_embedding = self.embedder.embed_query(query)
            chroma_results = self.vectorstore.search_communities(query_embedding, top_k=top_k)
            community_ids = [r["community_id"] for r in chroma_results]
            all_communities = self.repository.list_communities()
            id_to_community = {c["id"]: c for c in all_communities}
            top_communities = [id_to_community[cid] for cid in community_ids if cid in id_to_community]
        else:
            # Fallback: keyword matching
            communities = self.repository.list_communities()
            scored_communities = []
            query_terms = set(query.lower().split())
            for community in communities:
                summary = community.get("summary", "").lower()
                label = community.get("label", "").lower()
                score = 0
                for term in query_terms:
                    if term in summary:
                        score += 2
                    if term in label:
                        score += 1
                if score > 0:
                    scored_communities.append((community, score))
            scored_communities.sort(key=lambda x: (-x[1], x[0].get("id", 0)))
            top_communities = [c for c, _ in scored_communities[:top_k]]

        # Collect all entities in top communities
        all_entities = []
        for community in top_communities:
            entities = self.repository.list_entities_in_community(community["id"])
            all_entities.extend(entities)

        # LLM synthesis if provider is available
        synthesized_answer: str | None = None
        if self.provider and top_communities:
            summaries_text = "\n\n".join(
                f"Community: {c.get('label', '')}\n{c.get('summary', '')}"
                for c in top_communities
            )
            try:
                answer = self.provider.synthesize_global_answer(query, summaries_text)
                synthesized_answer = answer or None
            except Exception:
                pass

        return {
            "query": query,
            "communities": top_communities,
            "entities": all_entities,
            "count": len(top_communities),
            "synthesized_answer": synthesized_answer,
        }
