"""GraphRAGSearchProvider — registers GraphRAG as a search provider in unified_search."""

from __future__ import annotations

import logging
import uuid
from typing import Any

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from app.plugins.builtins.unified_search.contracts import BaseSearchProvider

logger = logging.getLogger(__name__)


class GraphRAGSearchProvider(BaseSearchProvider):
    supports_fts: bool = False
    supports_vector: bool = False
    supports_rag: bool = False
    supports_graph: bool = True
    """Search provider for GraphRAG entity relationships.

    Enables full-text and semantic search over relationship metadata
    and entity types in the knowledge graph.
    """

    entity_type = "graph_relationship"

    async def _search_fts_filtered(
        self,
        db: AsyncSession,
        tsquery: str,
        tenant_id: uuid.UUID,
        limit: int,
        visible_ids: set[uuid.UUID] | None,
    ) -> list[dict[str, Any]]:
        """Full-text search on relationship metadata and types."""
        if visible_ids is not None:
            sql = text(
                """
                SELECT r.*, ts_rank(
                    to_tsvector('pg_catalog.german',
                        coalesce(r.relationship_type, '') || ' ' ||
                        coalesce(r.source_type, '') || ' ' ||
                        coalesce(r.target_type, '') || ' ' ||
                        coalesce(r.meta::text, '')
                    ),
                    to_tsquery('pg_catalog.german', :q)
                ) AS rank
                FROM entity_relationships r
                WHERE r.tenant_id = :tid
                  AND r.deleted_at IS NULL
                  AND to_tsvector('pg_catalog.german',
                      coalesce(r.relationship_type, '') || ' ' ||
                      coalesce(r.source_type, '') || ' ' ||
                      coalesce(r.target_type, '') || ' ' ||
                      coalesce(r.meta::text, '')
                  ) @@ to_tsquery('pg_catalog.german', :q)
                  AND r.id = ANY(:visible_ids)
                ORDER BY rank DESC
                LIMIT :lim
                """
            )
            result = await db.execute(
                sql,
                {
                    "q": tsquery,
                    "tid": tenant_id,
                    "lim": limit,
                    "visible_ids": list(visible_ids),
                },
            )
        else:
            sql = text(
                """
                SELECT r.*, ts_rank(
                    to_tsvector('pg_catalog.german',
                        coalesce(r.relationship_type, '') || ' ' ||
                        coalesce(r.source_type, '') || ' ' ||
                        coalesce(r.target_type, '') || ' ' ||
                        coalesce(r.meta::text, '')
                    ),
                    to_tsquery('pg_catalog.german', :q)
                ) AS rank
                FROM entity_relationships r
                WHERE r.tenant_id = :tid
                  AND r.deleted_at IS NULL
                  AND to_tsvector('pg_catalog.german',
                      coalesce(r.relationship_type, '') || ' ' ||
                      coalesce(r.source_type, '') || ' ' ||
                      coalesce(r.target_type, '') || ' ' ||
                      coalesce(r.meta::text, '')
                  ) @@ to_tsquery('pg_catalog.german', :q)
                ORDER BY rank DESC
                LIMIT :lim
                """
            )
            result = await db.execute(
                sql,
                {"q": tsquery, "tid": tenant_id, "lim": limit},
            )
        rows = result.mappings().all()
        return [dict(r) for r in rows]

    async def _search_vector_filtered(
        self,
        db: AsyncSession,
        embedding: list[float],
        tenant_id: uuid.UUID,
        limit: int,
        visible_ids: set[uuid.UUID] | None,
    ) -> list[dict[str, Any]]:
        """Semantic search is not yet supported for graph relationships.

        Returns empty list — relationships are searched via FTS on metadata.
        """
        return []

    async def get_embedding_text(
        self, db: AsyncSession, entity_id: uuid.UUID, tenant_id: uuid.UUID
    ) -> str:
        """Get text for embedding generation."""
        sql = text(
            """
            SELECT relationship_type, source_type, source_id, target_type, target_id, metadata
            FROM entity_relationships
            WHERE id = :eid AND tenant_id = :tid
            """
        )
        result = await db.execute(sql, {"eid": entity_id, "tid": tenant_id})
        row = result.mappings().first()
        if not row:
            return ""
        parts = [
            row.get("relationship_type", ""),
            row.get("source_type", ""),
            str(row.get("source_id", "")),
            row.get("target_type", ""),
            str(row.get("target_id", "")),
            str(row.get("metadata", {})),
        ]
        return " ".join(str(p) for p in parts if p)

    async def search_graph(
        self,
        db: AsyncSession,
        query_analysis: dict[str, Any],
        tenant_id: uuid.UUID,
        limit: int,
        user_id: uuid.UUID | None = None,
        is_system_admin: bool = False,
    ) -> list[dict[str, Any]]:
        """Graph traversal search using BFS on entity_relationships.

        Starts from entities found by FTS on relationship metadata, then
        expands via BFS to find related entities up to depth 2.
        Returns related entities with relationship metadata.
        """
        normalized_query = query_analysis.get("normalized_query", "")
        semantic_terms = query_analysis.get("semantic_terms", [])
        tsquery_parts = [normalized_query] + semantic_terms
        tsquery = " & ".join(
            part.strip().replace(" ", " & ")
            for part in tsquery_parts
            if part and part.strip()
        )
        if not tsquery:
            tsquery = normalized_query
        if not tsquery:
            return []

        # Step 1: Find seed relationships via FTS
        seed_sql = text(
            """
            SELECT r.id, r.source_type, r.source_id, r.target_type, r.target_id,
                   r.relationship_type, r.metadata
            FROM entity_relationships r
            WHERE r.tenant_id = :tid
              AND r.deleted_at IS NULL
              AND to_tsvector('pg_catalog.german',
                  coalesce(r.relationship_type, '') || ' ' ||
                  coalesce(r.source_type, '') || ' ' ||
                  coalesce(r.target_type, '') || ' ' ||
                  coalesce(r.metadata::text, '')
              ) @@ to_tsquery('pg_catalog.german', :q)
            LIMIT :lim
            """
        )
        result = await db.execute(seed_sql, {"tid": tenant_id, "q": tsquery, "lim": limit})
        seed_rows = result.mappings().all()
        if not seed_rows:
            return []

        # Collect seed entity IDs (both source and target)
        seed_entities: dict[str, set[str]] = {}
        for row in seed_rows:
            for role in ("source", "target"):
                etype = row[f"{role}_type"]
                eid = str(row[f"{role}_id"])
                seed_entities.setdefault(etype, set()).add(eid)

        # Step 2: BFS traversal — find relationships connected to seed entities
        all_entity_ids: dict[str, set[str]] = {}
        for etype, eids in seed_entities.items():
            all_entity_ids.setdefault(etype, set()).update(eids)

        visited_rel_ids: set[str] = {str(r["id"]) for r in seed_rows}
        results: list[dict[str, Any]] = []

        # Add seed results first
        for row in seed_rows:
            results.append(self._relationship_to_result(dict(row)))

        # BFS: expand from seed entities (depth 1)
        if len(results) < limit:
            remaining = limit - len(results)
            bfs_sql = text(
                """
                SELECT r.id, r.source_type, r.source_id, r.target_type, r.target_id,
                       r.relationship_type, r.metadata
                FROM entity_relationships r
                WHERE r.tenant_id = :tid
                  AND r.deleted_at IS NULL
                  AND (
                    (r.source_type = :etype AND r.source_id = ANY(:eids))
                    OR (r.target_type = :etype AND r.target_id = ANY(:eids))
                  )
                  AND r.id <> ALL(:exclude_ids)
                LIMIT :lim
                """
            )
            for etype, eids in seed_entities.items():
                if len(results) >= limit:
                    break
                bfs_result = await db.execute(
                    bfs_sql,
                    {
                        "tid": tenant_id,
                        "etype": etype,
                        "eids": list(eids),
                        "exclude_ids": list(visited_rel_ids),
                        "lim": remaining,
                    },
                )
                bfs_rows = bfs_result.mappings().all()
                for row in bfs_rows:
                    rid = str(row["id"])
                    if rid not in visited_rel_ids:
                        visited_rel_ids.add(rid)
                        results.append(self._relationship_to_result(dict(row)))
                        if len(results) >= limit:
                            break

        return results[:limit]

    def _relationship_to_result(self, row: dict[str, Any]) -> dict[str, Any]:
        """Convert a relationship row to a search result dict."""
        return {
            "id": str(row.get("id", "")),
            "entity_type": self.entity_type,
            "entity_id": str(row.get("id", "")),
            "title": f"{row.get('source_type', '?')} --[{row.get('relationship_type', '?')}]--> {row.get('target_type', '?')}",
            "snippet": str(row.get("metadata", {})),
            "score": 0.0,
            "data": {
                "source_type": row.get("source_type"),
                "source_id": str(row.get("source_id", "")),
                "target_type": row.get("target_type"),
                "target_id": str(row.get("target_id", "")),
                "relationship_type": row.get("relationship_type"),
            },
        }

    def to_search_result(self, entity: object) -> dict[str, Any]:
        """Convert relationship to search result dict."""
        if isinstance(entity, dict):
            return {
                "entity_type": self.entity_type,
                "entity_id": str(entity.get("id", "")),
                "title": f"{entity.get('source_type', '?')} --[{entity.get('relationship_type', '?')}]--> {entity.get('target_type', '?')}",
                "snippet": str(entity.get("metadata", {})),
                "score": float(entity.get("rank", 0.0)),
                "data": {
                    "source_type": entity.get("source_type"),
                    "source_id": str(entity.get("source_id", "")),
                    "target_type": entity.get("target_type"),
                    "target_id": str(entity.get("target_id", "")),
                    "relationship_type": entity.get("relationship_type"),
                },
            }
        return {
            "entity_type": self.entity_type,
            "entity_id": str(getattr(entity, "id", "")),
            "title": f"{getattr(entity, 'source_type', '?')} --[{getattr(entity, 'relationship_type', '?')}]--> {getattr(entity, 'target_type', '?')}",
            "snippet": str(getattr(entity, "metadata", {})),
            "score": 0.0,
            "data": {},
        }