"""cortex/storage/queries/brain_graph_queries.py — Brain graph storage mixin."""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from cortex.storage._helpers import _now, _row_to_brain_graph, _row_to_brain_node, _row_to_brain_edge


class BrainGraphQueryMixin:
    """Mixin providing braingraphquery methods."""

    async def create_brain_graph(
        self, name: str, source_type: str, owner_id: str = "default", **kwargs
    ) -> BrainGraphIndex:
        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()
            async with conn.execute(
                """
                INSERT INTO brain_graph_index
                    (owner_id, name, source_type, source_uri, source_hash,
                     embedding_model_id, builder_prompt_ver, supersedes_id,
                     status, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                RETURNING *
                """,
                (
                    owner_id, name, source_type,
                    kwargs.get("source_uri"),
                    kwargs.get("source_hash"),
                    kwargs.get("embedding_model_id"),
                    kwargs.get("builder_prompt_ver"),
                    kwargs.get("supersedes_id"),
                    kwargs.get("status", "building"),
                    now,
                ),
            ) as cur:
                row = await cur.fetchone()
            await conn.commit()
            return _row_to_brain_graph(row)


    async def add_brain_graph_node(
        self, graph_id: str, node_type: str, label: str, **kwargs
    ) -> BrainGraphNode:
        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()
            async with conn.execute(
                """
                INSERT INTO brain_graph_node
                    (graph_id, node_type, label, content, compact_repr,
                     depth, parent_id, metadata_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                RETURNING *
                """,
                (
                    graph_id, node_type, label,
                    kwargs.get("content"),
                    kwargs.get("compact_repr"),
                    kwargs.get("depth", 0),
                    kwargs.get("parent_id"),
                    kwargs.get("metadata_json"),
                    now,
                ),
            ) as cur:
                row = await cur.fetchone()
            await conn.commit()

            # Update node count
            await conn.execute(
                "UPDATE brain_graph_index SET node_count = node_count + 1 WHERE id = ?",
                (graph_id,),
            )
            await conn.commit()
            return _row_to_brain_node(row)


    async def add_brain_graph_edge(
        self,
        graph_id: str,
        source_id: str,
        target_id: str,
        edge_type: str,
        **kwargs,
    ) -> BrainGraphEdge:
        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()
            async with conn.execute(
                """
                INSERT INTO brain_graph_edge
                    (graph_id, source_id, target_id, edge_type,
                     weight, metadata_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                RETURNING *
                """,
                (
                    graph_id, source_id, target_id, edge_type,
                    kwargs.get("weight", 1.0),
                    kwargs.get("metadata_json"),
                    now,
                ),
            ) as cur:
                row = await cur.fetchone()
            await conn.commit()

            await conn.execute(
                "UPDATE brain_graph_index SET edge_count = edge_count + 1 WHERE id = ?",
                (graph_id,),
            )
            await conn.commit()
            return _row_to_brain_edge(row)


    async def get_brain_graph_nodes(
        self, graph_id: str, node_type: Optional[str] = None
    ) -> List[BrainGraphNode]:
        conn = await self._get_conn()
        if node_type:
            async with conn.execute(
                "SELECT * FROM brain_graph_node WHERE graph_id = ? AND node_type = ?",
                (graph_id, node_type),
            ) as cur:
                rows = await cur.fetchall()
        else:
            async with conn.execute(
                "SELECT * FROM brain_graph_node WHERE graph_id = ?",
                (graph_id,),
            ) as cur:
                rows = await cur.fetchall()
        return [_row_to_brain_node(r) for r in rows]


    async def get_brain_graph(self, graph_id: str) -> Optional[BrainGraphIndex]:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM brain_graph_index WHERE id = ?", (graph_id,)
        ) as cur:
            row = await cur.fetchone()
        return _row_to_brain_graph(row) if row else None


    async def find_brain_graph_by_hash(
        self, source_hash: str, owner_id: str = "default"
    ) -> Optional[BrainGraphIndex]:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM brain_graph_index WHERE owner_id = ? AND source_hash = ? AND status = 'active'",
            (owner_id, source_hash),
        ) as cur:
            row = await cur.fetchone()
        return _row_to_brain_graph(row) if row else None


    async def list_brain_graphs(
        self, owner_id: str = "default"
    ) -> List[BrainGraphIndex]:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM brain_graph_index WHERE owner_id = ? ORDER BY created_at DESC",
            (owner_id,),
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_brain_graph(r) for r in rows]


    async def update_brain_graph(self, graph_id: str, **kwargs) -> BrainGraphIndex:
        async with self._write_lock:
            conn = await self._get_conn()
            allowed = {
                "name", "status", "source_hash", "source_uri", "embedding_model_id",
                "builder_prompt_ver", "supersedes_id", "node_count", "edge_count",
                "last_indexed", "index_config",
            }
            updates = {k: v for k, v in kwargs.items() if k in allowed}
            if not updates:
                row_obj = await self.get_brain_graph(graph_id)
                assert row_obj is not None
                return row_obj
            cols = ", ".join(f"{k} = ?" for k in updates)
            vals = list(updates.values()) + [graph_id]
            await conn.execute(
                f"UPDATE brain_graph_index SET {cols} WHERE id = ?", vals
            )
            await conn.commit()
            row_obj = await self.get_brain_graph(graph_id)
            assert row_obj is not None
            return row_obj


    async def delete_brain_graph(self, graph_id: str) -> None:
        async with self._write_lock:
            conn = await self._get_conn()
            # Cascade delete nodes and edges first
            await conn.execute(
                "DELETE FROM brain_graph_edge WHERE graph_id = ?", (graph_id,)
            )
            await conn.execute(
                "DELETE FROM brain_graph_node WHERE graph_id = ?", (graph_id,)
            )
            await conn.execute(
                "DELETE FROM brain_graph_index WHERE id = ?", (graph_id,)
            )
            await conn.commit()


    async def get_brain_graph_node(self, node_id: str) -> Optional[BrainGraphNode]:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM brain_graph_node WHERE id = ?", (node_id,)
        ) as cur:
            row = await cur.fetchone()
        return _row_to_brain_node(row) if row else None


    async def update_brain_graph_node(self, node_id: str, **kwargs) -> BrainGraphNode:
        async with self._write_lock:
            conn = await self._get_conn()
            allowed = {"node_type", "label", "content", "compact_repr", "depth", "parent_id", "metadata_json"}
            updates = {k: v for k, v in kwargs.items() if k in allowed}
            if updates:
                cols = ", ".join(f"{k} = ?" for k in updates)
                vals = list(updates.values()) + [node_id]
                await conn.execute(
                    f"UPDATE brain_graph_node SET {cols} WHERE id = ?", vals
                )
                await conn.commit()
            row_obj = await self.get_brain_graph_node(node_id)
            assert row_obj is not None
            return row_obj


    async def delete_brain_graph_node(self, node_id: str) -> None:
        async with self._write_lock:
            conn = await self._get_conn()
            # Also delete edges referencing this node
            await conn.execute(
                "DELETE FROM brain_graph_edge WHERE source_id = ? OR target_id = ?",
                (node_id, node_id),
            )
            await conn.execute(
                "DELETE FROM brain_graph_node WHERE id = ?", (node_id,)
            )
            await conn.commit()


    async def get_brain_graph_edges(self, graph_id: str) -> List[BrainGraphEdge]:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM brain_graph_edge WHERE graph_id = ?", (graph_id,)
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_brain_edge(r) for r in rows]


    async def get_brain_graph_edge(self, edge_id: str) -> Optional[BrainGraphEdge]:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM brain_graph_edge WHERE id = ?", (edge_id,)
        ) as cur:
            row = await cur.fetchone()
        return _row_to_brain_edge(row) if row else None


    async def update_brain_graph_edge(self, edge_id: str, **kwargs) -> BrainGraphEdge:
        async with self._write_lock:
            conn = await self._get_conn()
            allowed = {"edge_type", "weight", "metadata_json"}
            updates = {k: v for k, v in kwargs.items() if k in allowed}
            if updates:
                cols = ", ".join(f"{k} = ?" for k in updates)
                vals = list(updates.values()) + [edge_id]
                await conn.execute(
                    f"UPDATE brain_graph_edge SET {cols} WHERE id = ?", vals
                )
                await conn.commit()
            row_obj = await self.get_brain_graph_edge(edge_id)
            assert row_obj is not None
            return row_obj


    async def delete_brain_graph_edge(self, edge_id: str) -> None:
        async with self._write_lock:
            conn = await self._get_conn()
            await conn.execute(
                "DELETE FROM brain_graph_edge WHERE id = ?", (edge_id,)
            )
            await conn.commit()


    async def list_brain_nodes(
        self,
        node_type: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
        owner_id: Optional[str] = None,
    ) -> List[BrainGraphNode]:
        """List brain graph nodes across all graphs, with optional type filter.

        Joins brain_graph_node with brain_graph_index so we can filter by
        owner_id without requiring the caller to know individual graph IDs.
        """
        conn = await self._get_conn()
        params: list = []
        where_clauses: list = []

        if owner_id is not None:
            where_clauses.append("bgi.owner_id = ?")
            params.append(owner_id)
        if node_type is not None:
            where_clauses.append("bgn.node_type = ?")
            params.append(node_type)

        where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
        sql = (
            "SELECT bgn.* FROM brain_graph_node bgn "
            "JOIN brain_graph_index bgi ON bgn.graph_id = bgi.id "
            f"{where_sql} ORDER BY bgn.created_at DESC LIMIT ? OFFSET ?"
        )
        params.extend([limit, offset])
        async with conn.execute(sql, params) as cur:
            rows = await cur.fetchall()
        return [_row_to_brain_node(r) for r in rows]


    async def count_brain_nodes(
        self,
        node_type: Optional[str] = None,
        owner_id: Optional[str] = None,
    ) -> int:
        """Return total count of brain graph nodes (optionally filtered)."""
        conn = await self._get_conn()
        params: list = []
        where_clauses: list = []

        if owner_id is not None:
            where_clauses.append("bgi.owner_id = ?")
            params.append(owner_id)
        if node_type is not None:
            where_clauses.append("bgn.node_type = ?")
            params.append(node_type)

        where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
        sql = (
            "SELECT COUNT(*) FROM brain_graph_node bgn "
            "JOIN brain_graph_index bgi ON bgn.graph_id = bgi.id "
            f"{where_sql}"
        )
        async with conn.execute(sql, params) as cur:
            row = await cur.fetchone()
        return row[0] if row else 0


    async def get_brain_node(self, node_id: str) -> Optional[BrainGraphNode]:
        """Return a single brain graph node by ID (alias for get_brain_graph_node)."""
        return await self.get_brain_graph_node(node_id)


    async def delete_brain_node(self, node_id: str) -> None:
        """Delete a brain graph node and its connected edges (alias for delete_brain_graph_node)."""
        await self.delete_brain_graph_node(node_id)


    async def get_brain_stats(self, owner_id: Optional[str] = None) -> dict:
        """Return aggregate brain graph statistics for the given owner."""
        conn = await self._get_conn()
        params_owner: list = [owner_id] if owner_id is not None else []
        owner_filter = "WHERE owner_id = ?" if owner_id is not None else ""

        # Total nodes across owner's graphs
        total_nodes = 0
        total_edges = 0
        node_types: dict = {}
        last_build: Optional[str] = None

        try:
            # Total nodes
            sql_nodes = (
                "SELECT COUNT(*) FROM brain_graph_node bgn "
                "JOIN brain_graph_index bgi ON bgn.graph_id = bgi.id "
                + ("WHERE bgi.owner_id = ?" if owner_id else "")
            )
            async with conn.execute(sql_nodes, params_owner) as cur:
                row = await cur.fetchone()
                total_nodes = row[0] if row else 0

            # Total edges
            sql_edges = (
                "SELECT COUNT(*) FROM brain_graph_edge bge "
                "JOIN brain_graph_index bgi ON bge.graph_id = bgi.id "
                + ("WHERE bgi.owner_id = ?" if owner_id else "")
            )
            async with conn.execute(sql_edges, params_owner) as cur:
                row = await cur.fetchone()
                total_edges = row[0] if row else 0

            # Node types breakdown
            sql_types = (
                "SELECT bgn.node_type, COUNT(*) FROM brain_graph_node bgn "
                "JOIN brain_graph_index bgi ON bgn.graph_id = bgi.id "
                + ("WHERE bgi.owner_id = ? " if owner_id else "")
                + "GROUP BY bgn.node_type"
            )
            async with conn.execute(sql_types, params_owner) as cur:
                async for type_row in cur:
                    node_types[type_row[0] or "unknown"] = type_row[1]

            # Last build timestamp
            sql_last = (
                "SELECT created_at FROM brain_graph_index "
                + owner_filter
                + " ORDER BY created_at DESC LIMIT 1"
            )
            async with conn.execute(sql_last, params_owner) as cur:
                row = await cur.fetchone()
                last_build = row[0] if row else None
        except Exception as e:
            import logging
            logging.getLogger(__name__).warning("get_brain_stats: error collecting stats: %s", e)

        return {
            "total_nodes": total_nodes,
            "total_edges": total_edges,
            "node_types": node_types,
            "last_build": last_build,
        }


