"""cortex/storage/queries/graph_queries.py — Graph visualization storage mixin."""

from __future__ import annotations

import collections
from typing import Any, Dict, List, Optional


class GraphQueryMixin:
    """Mixin providing graphquery methods."""

    async def get_entity_graph(
        self,
        owner_id: str,
        center_entity_id: Optional[str] = None,
        max_depth: int = 2,
        max_nodes: int = 50,
    ) -> Dict[str, Any]:
        """
        Get entity graph for force-directed visualization.
        
        Returns nodes (entities) + edges (relationships derived from claims).
        An edge exists when two entities appear in the same claim
        (one as subject, other as object).
        Edge label = predicate from the claim.
        
        If center_entity_id: start from that entity, expand outward to max_depth.
        If no center: return the most connected entities.
        """
        conn = await self._get_conn()
        
        # Color mapping for entity types
        entity_type_colors = {
            "person": "#4A90D9",
            "place": "#50C878",
            "organization": "#F5A623",
            "project": "#9013FE",
            "product": "#D0021B",
            "concept": "#B8E986",
            "animal": "#F8E71C",
            "unknown": "#9B9B9B",
        }
        
        nodes = []
        edges = []
        entity_ids_found = set()
        
        if center_entity_id:
            # BFS to find entities within max_depth
            # Start with the center entity
            current_level = {center_entity_id}
            visited = {center_entity_id}
            
            for depth in range(max_depth):
                if not current_level or len(entity_ids_found) >= max_nodes:
                    break
                    
                # Find all claims involving current entities
                placeholders = ",".join("?" * len(current_level))
                level_list = list(current_level)
                async with conn.execute(
                    f"""
                    SELECT DISTINCT 
                           COALESCE(sc.subject_entity_id, '') as subj,
                           COALESCE(sc.object_entity_id, '') as obj,
                           sc.predicate,
                           sc.id as claim_id
                    FROM semantic_claim sc
                    WHERE (sc.subject_entity_id IN ({placeholders})
                           OR sc.object_entity_id IN ({placeholders}))
                      AND sc.status = 'active'
                      AND (sc.subject_entity_id IS NOT NULL 
                           AND sc.subject_entity_id != ''
                           AND sc.object_entity_id IS NOT NULL 
                           AND sc.object_entity_id != '')
                    """,
                    level_list + level_list,  # Same values for both subj and obj IN clauses
                ) as cur:
                    rows = await cur.fetchall()
                
                next_level = set()
                for row in rows:
                    subj = row["subj"]
                    obj = row["obj"]
                    predicate = row["predicate"]
                    claim_id = row["claim_id"]
                    
                    if subj and subj not in visited:
                        next_level.add(subj)
                        visited.add(subj)
                    if obj and obj not in visited:
                        next_level.add(obj)
                        visited.add(obj)
                    
                    # Add edge
                    if subj and obj and subj != obj:
                        # Check if edge already exists
                        edge_key = tuple(sorted([subj, obj]))
                        edge_exists = False
                        for e in edges:
                            if e["source"] == subj and e["target"] == obj:
                                e["weight"] += 1
                                edge_exists = True
                                break
                        if not edge_exists:
                            edges.append({
                                "source": subj,
                                "target": obj,
                                "label": predicate,
                                "weight": 1,
                            })
                
                current_level = next_level
            
            entity_ids_found = visited
        else:
            # No center entity - return most connected entities
            # Find entities with most relationships
            async with conn.execute(
                """
                SELECT 
                    e.id,
                    e.canonical_name,
                    e.entity_type,
                    e.mention_count,
                    COUNT(DISTINCT sc.id) as claim_count
                FROM entity e
                LEFT JOIN semantic_claim sc 
                    ON (sc.subject_entity_id = e.id OR sc.object_entity_id = e.id)
                    AND sc.status = 'active'
                WHERE e.owner_id = ?
                GROUP BY e.id
                ORDER BY claim_count DESC, e.mention_count DESC
                LIMIT ?
                """,
                (owner_id, max_nodes),
            ) as cur:
                rows = await cur.fetchall()
            
            entity_ids_found = {row["id"] for row in rows}
            
            # Get edges between these entities
            if entity_ids_found:
                placeholders = ",".join("?" * len(entity_ids_found))
                entity_list = list(entity_ids_found)
                async with conn.execute(
                    f"""
                    SELECT 
                        sc.subject_entity_id,
                        sc.object_entity_id,
                        sc.predicate,
                        COUNT(*) as weight
                    FROM semantic_claim sc
                    WHERE sc.subject_entity_id IN ({placeholders})
                      AND sc.object_entity_id IN ({placeholders})
                      AND sc.status = 'active'
                      AND sc.subject_entity_id IS NOT NULL
                      AND sc.object_entity_id IS NOT NULL
                    GROUP BY sc.subject_entity_id, sc.object_entity_id, sc.predicate
                    """,
                    entity_list + entity_list,
                ) as cur:
                    edge_rows = await cur.fetchall()
                
                for row in edge_rows:
                    if row["subject_entity_id"] and row["object_entity_id"]:
                        edges.append({
                            "source": row["subject_entity_id"],
                            "target": row["object_entity_id"],
                            "label": row["predicate"],
                            "weight": row["weight"],
                        })
        
        # Fetch entity details for nodes
        if entity_ids_found:
            placeholders = ",".join("?" * len(entity_ids_found))
            async with conn.execute(
                f"""
                SELECT id, canonical_name, entity_type, mention_count
                FROM entity
                WHERE id IN ({placeholders})
                """,
                list(entity_ids_found),
            ) as cur:
                entity_rows = await cur.fetchall()
            
            for row in entity_rows:
                nodes.append({
                    "id": row["id"],
                    "label": row["canonical_name"],
                    "type": row["entity_type"],
                    "size": min(row["mention_count"] or 1, 30),  # Cap size at 30
                    "color": entity_type_colors.get(row["entity_type"], "#9B9B9B"),
                })
        
        return {
            "nodes": nodes,
            "edges": edges,
            "total_entities": len(nodes),
            "total_relationships": len(edges),
        }


    async def get_entity_neighbors(
        self,
        entity_id: str,
        owner_id: str,
    ) -> Dict[str, Any]:
        """
        Get direct neighbors only (depth=1) for an entity.
        
        Returns nodes + edges representing immediate connections.
        """
        conn = await self._get_conn()
        
        # Color mapping for entity types
        entity_type_colors = {
            "person": "#4A90D9",
            "place": "#50C878",
            "organization": "#F5A623",
            "project": "#9013FE",
            "product": "#D0021B",
            "concept": "#B8E986",
            "animal": "#F8E71C",
            "unknown": "#9B9B9B",
        }
        
        # Get the center entity
        async with conn.execute(
            "SELECT id, canonical_name, entity_type, mention_count FROM entity WHERE id = ?",
            (entity_id,),
        ) as cur:
            center_row = await cur.fetchone()
        
        if not center_row:
            return {
                "nodes": [],
                "edges": [],
                "total_entities": 0,
                "total_relationships": 0,
            }
        
        nodes = [{
            "id": center_row["id"],
            "label": center_row["canonical_name"],
            "type": center_row["entity_type"],
            "size": min(center_row["mention_count"] or 1, 30),
            "color": entity_type_colors.get(center_row["entity_type"], "#9B9B9B"),
        }]
        
        # Get direct neighbors via claims
        async with conn.execute(
            """
            SELECT 
                e.id,
                e.canonical_name,
                e.entity_type,
                e.mention_count,
                sc.predicate,
                CASE 
                    WHEN sc.subject_entity_id = ? THEN 'outgoing'
                    ELSE 'incoming'
                END as direction
            FROM semantic_claim sc
            JOIN entity e ON (
                CASE 
                    WHEN sc.subject_entity_id = ? THEN sc.object_entity_id
                    ELSE sc.subject_entity_id
                END = e.id
            )
            WHERE (sc.subject_entity_id = ? OR sc.object_entity_id = ?)
              AND sc.status = 'active'
              AND (
                  (sc.subject_entity_id = ? AND sc.object_entity_id IS NOT NULL)
                  OR (sc.object_entity_id = ? AND sc.subject_entity_id IS NOT NULL)
              )
            GROUP BY e.id, sc.predicate
            """,
            (entity_id, entity_id, entity_id, entity_id, entity_id, entity_id),
        ) as cur:
            neighbor_rows = await cur.fetchall()
        
        edges = []
        neighbor_ids = set()
        
        for row in neighbor_rows:
            if row["id"] not in neighbor_ids:
                neighbor_ids.add(row["id"])
                nodes.append({
                    "id": row["id"],
                    "label": row["canonical_name"],
                    "type": row["entity_type"],
                    "size": min(row["mention_count"] or 1, 30),
                    "color": entity_type_colors.get(row["entity_type"], "#9B9B9B"),
                })
            
            edges.append({
                "source": entity_id if row["direction"] == "outgoing" else row["id"],
                "target": row["id"] if row["direction"] == "outgoing" else entity_id,
                "label": row["predicate"],
                "weight": 1,
            })
        
        return {
            "nodes": nodes,
            "edges": edges,
            "total_entities": len(nodes),
            "total_relationships": len(edges),
        }


    async def find_shortest_path(
        self,
        source_id: str,
        target_id: str,
        owner_id: str,
        max_hops: int = 5,
    ) -> List[Dict[str, Any]]:
        """
        Find shortest path between two entities using BFS.
        Max hops defaults to 5.
        
        Returns list of path segments with entity and relationship info.
        """
        conn = await self._get_conn()
        
        if source_id == target_id:
            return [{"entity_id": source_id, "predicate": None, "step": 0}]
        
        # BFS (using deque for O(1) popleft instead of O(n) list.pop(0))
        visited = {source_id}
        queue = collections.deque([(source_id, [])])  # (current_entity, path_so_far)
        
        while queue:
            current, path = queue.popleft()
            
            if len(path) >= max_hops:
                continue
            
            # Find all direct neighbors
            async with conn.execute(
                """
                SELECT 
                    CASE 
                        WHEN sc.subject_entity_id = ? THEN sc.object_entity_id
                        ELSE sc.subject_entity_id
                    END as neighbor_id,
                    sc.predicate,
                    sc.subject_entity_id,
                    sc.object_entity_id
                FROM semantic_claim sc
                WHERE (sc.subject_entity_id = ? OR sc.object_entity_id = ?)
                  AND sc.status = 'active'
                  AND (
                      (sc.subject_entity_id = ? AND sc.object_entity_id IS NOT NULL)
                      OR (sc.object_entity_id = ? AND sc.subject_entity_id IS NOT NULL)
                  )
                """,
                (current, current, current, current, current),
            ) as cur:
                rows = await cur.fetchall()
            
            for row in rows:
                neighbor_id = row["neighbor_id"]
                if not neighbor_id or neighbor_id == current:
                    continue
                
                if neighbor_id == target_id:
                    # Found the target - reconstruct path
                    final_path = path + [{
                        "entity_id": neighbor_id,
                        "predicate": row["predicate"],
                        "step": len(path) + 1,
                        "from_entity": row["subject_entity_id"],
                        "to_entity": row["object_entity_id"],
                    }]
                    return final_path
                
                if neighbor_id not in visited:
                    visited.add(neighbor_id)
                    new_path = path + [{
                        "entity_id": neighbor_id,
                        "predicate": row["predicate"],
                        "step": len(path) + 1,
                        "from_entity": row["subject_entity_id"],
                        "to_entity": row["object_entity_id"],
                    }]
                    queue.append((neighbor_id, new_path))
        
        # No path found within max_hops
        return []

