"""LlamaIndex PropertyGraphStore backed by Apache AGE + pgvector.

Implements the modern LlamaIndex PropertyGraphStore interface with labeled
nodes, typed relations, and built-in vector query support.

Architecture: Graph structure lives in AGE (vertices/edges), embeddings
live in a companion pgvector table. This follows Piggie's principle:
graph for structure, SQL for measures/embeddings.
"""

from __future__ import annotations

from typing import Any

from piggie._util import validate_label as _validate_label
from piggie.exceptions import IntegrationError
from piggie.types import Edge, Vertex

# ---------------------------------------------------------------------------
# Conditional imports — gracefully degrade when llama_index is absent
# ---------------------------------------------------------------------------

try:
    from llama_index.core.graph_stores.types import PropertyGraphStore
    from llama_index.core.schema import BaseNode

    _HAS_LLAMAINDEX = True
except ImportError:
    _HAS_LLAMAINDEX = False

    class PropertyGraphStore:  # type: ignore[no-redef]
        """Stub when llama_index is not installed."""
        supports_structured_queries: bool = False
        supports_vector_queries: bool = False

    class BaseNode:  # type: ignore[no-redef]
        pass

# Lazy-loaded types — only used when llama_index is actually installed
_EntityNode: type | None = None
_ChunkNode: type | None = None
_Relation: type | None = None
_LabelledNode: type | None = None
_VectorStoreQuery: type | None = None


def _load_llamaindex_types() -> None:
    """Import LlamaIndex types on first use."""
    global _EntityNode, _ChunkNode, _Relation, _LabelledNode, _VectorStoreQuery
    if _EntityNode is not None:
        return
    try:
        from llama_index.core.graph_stores.types import (
            ChunkNode,
            EntityNode,
            LabelledNode,
            Relation,
        )
        from llama_index.core.vector_stores.types import VectorStoreQuery
        _EntityNode = EntityNode
        _ChunkNode = ChunkNode
        _Relation = Relation
        _LabelledNode = LabelledNode
        _VectorStoreQuery = VectorStoreQuery
    except ImportError as exc:
        raise IntegrationError(
            "llama-index-core >= 0.10 is required for PropertyGraphStore. "
            "Install with: pip install 'llama-index-core>=0.10'"
        ) from exc


def _vertex_to_labelled_node(vertex: Vertex) -> Any:
    """Convert a Piggie Vertex to a LlamaIndex EntityNode or ChunkNode."""
    _load_llamaindex_types()
    assert _ChunkNode is not None
    assert _EntityNode is not None
    props = dict(vertex.properties)
    node_id = str(props.pop("id", vertex.id))

    if "text" in props:
        text = str(props.pop("text"))
        return _ChunkNode(
            id_=node_id,
            label=vertex.label,
            text=text,
            properties=props,
        )
    else:
        name = str(props.pop("name", node_id))
        return _EntityNode(
            id_=node_id,
            label=vertex.label,
            name=name,
            properties=props,
        )


class PiggiePropertyGraphStore(PropertyGraphStore):
    """LlamaIndex PropertyGraphStore backed by Apache AGE + pgvector.

    Graph structure is stored in AGE (vertices and edges). Node embeddings
    are stored in a companion pgvector table for vector similarity queries.

    Usage::

        from piggie._integrations.llamaindex_property import PiggiePropertyGraphStore
        store = PiggiePropertyGraphStore(driver, graph="my_graph")
    """

    supports_structured_queries: bool = True
    supports_vector_queries: bool = True

    def __init__(
        self,
        driver: Any,
        graph: str = "default",
        *,
        embedding_table: str | None = None,
        embedding_dimensions: int | None = None,
        create_tables: bool = True,
    ) -> None:
        self._driver = driver
        self._graph = graph
        self._embedding_table = embedding_table or f"{graph}_node_embeddings"
        self._embedding_dimensions = embedding_dimensions
        self._schema_cache: dict[str, Any] | None = None

        if create_tables:
            self._ensure_embedding_table()

    def _ensure_embedding_table(self) -> None:
        """Create the companion pgvector embedding table if it doesn't exist."""
        dim_spec = f"({self._embedding_dimensions})" if self._embedding_dimensions else ""
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self._embedding_table} ("
            f"node_id TEXT PRIMARY KEY, "
            f"node_label TEXT NOT NULL, "
            f"text TEXT, "
            f"embedding vector{dim_spec}, "
            f"properties JSONB DEFAULT '{{}}'"
            f")"
        )
        try:
            self._driver.sql(sql)
            self._driver.commit()
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(
                f"Failed to create embedding table: {exc}"
            ) from exc

    # ------------------------------------------------------------------
    # Node operations
    # ------------------------------------------------------------------

    def upsert_nodes(self, nodes: list[Any]) -> None:
        """Upsert nodes into AGE and optionally into the embedding table.

        Accepts LlamaIndex EntityNode and ChunkNode instances.
        """
        _load_llamaindex_types()
        try:
            for node in nodes:
                node_id = node.id_
                label = _validate_label(getattr(node, "label", "Entity") or "Entity")
                props = dict(getattr(node, "properties", {}))

                # Add name or text depending on node type
                if _EntityNode and isinstance(node, _EntityNode):
                    name = getattr(node, "name", node_id)
                    props["name"] = name
                elif _ChunkNode and isinstance(node, _ChunkNode):
                    text = getattr(node, "text", "")
                    props["text"] = text

                # Phase 1: MERGE vertex in AGE
                prop_keys = list(props.keys())
                prop_vals = list(props.values())
                set_parts = ", ".join(
                    f"n.{_validate_label(k)} = %s" for k in prop_keys
                )
                q = f"MERGE (n:{label} {{id: %s}})"
                cypher_params: tuple[Any, ...] = (node_id,)
                if set_parts:
                    q += f" SET {set_parts}"
                    cypher_params = (node_id, *prop_vals)
                self._driver.cypher(q, params=cypher_params, graph=self._graph)

                # Phase 2: UPSERT into embedding table if embedding present
                embedding = getattr(node, "embedding", None)
                if embedding is not None:
                    import json

                    text_val = props.get("text") or props.get("name") or ""
                    vec_str = "[" + ",".join(str(v) for v in embedding) + "]"
                    props_json = json.dumps(props)
                    sql = (
                        f"INSERT INTO {self._embedding_table} "
                        f"(node_id, node_label, text, embedding, properties) "
                        f"VALUES (%s, %s, %s, %s::vector, %s::jsonb) "
                        f"ON CONFLICT (node_id) DO UPDATE SET "
                        f"node_label = EXCLUDED.node_label, "
                        f"text = EXCLUDED.text, "
                        f"embedding = EXCLUDED.embedding, "
                        f"properties = EXCLUDED.properties"
                    )
                    self._driver.sql(
                        sql, (node_id, label, str(text_val), vec_str, props_json)
                    )

            self._driver.commit()
        except IntegrationError:
            self._driver.rollback()
            raise
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"upsert_nodes failed: {exc}") from exc

    def upsert_relations(self, relations: list[Any]) -> None:
        """Upsert edges into AGE."""
        _load_llamaindex_types()
        try:
            for rel in relations:
                source_id = rel.source_id
                target_id = rel.target_id
                rel_label = _validate_label(rel.label)

                props = dict(getattr(rel, "properties", {}))
                prop_keys = list(props.keys())
                prop_vals = list(props.values())

                # Build MERGE for source → target edge
                source_label = _validate_label(
                    getattr(rel, "source_label", "Entity") or "Entity"
                )
                target_label = _validate_label(
                    getattr(rel, "target_label", "Entity") or "Entity"
                )

                q = (
                    f"MATCH (s:{source_label} {{id: %s}}), "
                    f"(t:{target_label} {{id: %s}}) "
                    f"MERGE (s)-[r:{rel_label}]->(t)"
                )
                cypher_params: tuple[Any, ...] = (source_id, target_id)

                if prop_keys:
                    set_parts = ", ".join(
                        f"r.{_validate_label(k)} = %s" for k in prop_keys
                    )
                    q += f" SET {set_parts}"
                    cypher_params = (source_id, target_id, *prop_vals)

                self._driver.cypher(q, params=cypher_params, graph=self._graph)

            self._driver.commit()
        except IntegrationError:
            self._driver.rollback()
            raise
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"upsert_relations failed: {exc}") from exc

    # ------------------------------------------------------------------
    # Query operations
    # ------------------------------------------------------------------

    def get(
        self,
        properties: dict[str, Any] | None = None,
        ids: list[str] | None = None,
    ) -> list[Any]:
        """Get nodes by property filter or ID list."""
        _load_llamaindex_types()
        try:
            if ids:
                # Fetch by ID list — one query per ID since AGE doesn't support IN
                results = []
                for nid in ids:
                    q = "MATCH (n {id: %s}) RETURN n"
                    result = self._driver.cypher(
                        q, params=(nid,), graph=self._graph
                    )
                    for record in result:
                        v = record[0]
                        if isinstance(v, Vertex):
                            results.append(_vertex_to_labelled_node(v))
                return results

            if properties:
                # Build property filter
                keys = list(properties.keys())
                vals = list(properties.values())
                where_parts = " AND ".join(f"n.{_validate_label(k)} = %s" for k in keys)
                q = f"MATCH (n) WHERE {where_parts} RETURN n"
                result = self._driver.cypher(
                    q, params=tuple(vals), graph=self._graph
                )
                return [
                    _vertex_to_labelled_node(r[0])
                    for r in result
                    if isinstance(r[0], Vertex)
                ]

            # No filters — return all
            q = "MATCH (n) RETURN n"
            result = self._driver.cypher(q, graph=self._graph)
            return [
                _vertex_to_labelled_node(r[0])
                for r in result
                if isinstance(r[0], Vertex)
            ]

        except IntegrationError:
            raise
        except Exception as exc:
            raise IntegrationError(f"get failed: {exc}") from exc

    def get_triplets(
        self,
        entity_names: list[str] | None = None,
        relation_names: list[str] | None = None,
        properties: dict[str, Any] | None = None,
        ids: list[str] | None = None,
    ) -> list[Any]:
        """Get triplets (subject, relation, object) with optional filters.

        Returns list of (EntityNode, Relation, EntityNode) tuples.
        """
        _load_llamaindex_types()
        triplets: list[Any] = []

        try:
            if entity_names:
                # Filter by source entity name
                for name in entity_names:
                    q = (
                        "MATCH (s)-[r]->(o) "
                        "WHERE s.name = %s OR o.name = %s "
                        "RETURN s, r, o"
                    )
                    result = self._driver.cypher(
                        q, params=(name, name), graph=self._graph
                    )
                    for record in result:
                        t = self._record_to_triplet(record)
                        if t is not None:
                            triplets.append(t)
                return triplets

            if relation_names:
                for rel_name in relation_names:
                    _validate_label(rel_name)
                    q = f"MATCH (s)-[r:{rel_name}]->(o) RETURN s, r, o"
                    result = self._driver.cypher(q, graph=self._graph)
                    for record in result:
                        t = self._record_to_triplet(record)
                        if t is not None:
                            triplets.append(t)
                return triplets

            if ids:
                for nid in ids:
                    q = "MATCH (s {id: %s})-[r]->(o) RETURN s, r, o"
                    result = self._driver.cypher(
                        q, params=(nid,), graph=self._graph
                    )
                    for record in result:
                        t = self._record_to_triplet(record)
                        if t is not None:
                            triplets.append(t)
                return triplets

            # No filters — all triplets
            q = "MATCH (s)-[r]->(o) RETURN s, r, o"
            result = self._driver.cypher(q, graph=self._graph)
            for record in result:
                t = self._record_to_triplet(record)
                if t is not None:
                    triplets.append(t)
            return triplets

        except IntegrationError:
            raise
        except Exception as exc:
            raise IntegrationError(f"get_triplets failed: {exc}") from exc

    def _record_to_triplet(self, record: Any) -> Any | None:
        """Convert a (Vertex, Edge, Vertex) record to a LlamaIndex triplet."""
        assert _Relation is not None
        s, r, o = record[0], record[1], record[2]
        if not isinstance(s, Vertex) or not isinstance(o, Vertex):
            return None
        if not isinstance(r, Edge):
            return None

        source_node = _vertex_to_labelled_node(s)
        target_node = _vertex_to_labelled_node(o)

        rel = _Relation(
            source_id=source_node.id_,
            target_id=target_node.id_,
            label=r.label,
            properties=dict(r.properties),
        )
        return (source_node, rel, target_node)

    def get_rel_map(
        self,
        graph_nodes: list[Any] | None = None,
        depth: int = 2,
        limit: int = 30,
        ignore_rels: list[str] | None = None,
    ) -> list[Any]:
        """Get relationship map via VLE traversal from given nodes.

        Returns list of (source, relation, target) triplets.
        """
        _load_llamaindex_types()
        triplets: list[Any] = []
        ignore_set = set(ignore_rels or [])

        try:
            if graph_nodes:
                for node in graph_nodes:
                    node_id = node.id_
                    q = (
                        f"MATCH (s {{id: %s}})-[r*1..{depth}]->(o) "
                        f"RETURN s, r, o LIMIT {limit}"
                    )
                    result = self._driver.cypher(
                        q, params=(node_id,), graph=self._graph
                    )
                    for record in result:
                        s, r_val, o = record[0], record[1], record[2]
                        if not isinstance(s, Vertex) or not isinstance(o, Vertex):
                            continue

                        # r_val is a list of edges for VLE
                        edges = r_val if isinstance(r_val, list) else [r_val]
                        for edge in edges:
                            if not isinstance(edge, Edge):
                                continue
                            if edge.label in ignore_set:
                                continue
                            assert _Relation is not None
                            source_node = _vertex_to_labelled_node(s)
                            target_node = _vertex_to_labelled_node(o)
                            rel = _Relation(
                                source_id=source_node.id_,
                                target_id=target_node.id_,
                                label=edge.label,
                                properties=dict(edge.properties),
                            )
                            triplets.append((source_node, rel, target_node))
            else:
                # No starting nodes — return all direct relationships
                q = "MATCH (s)-[r]->(o) RETURN s, r, o LIMIT %s"
                result = self._driver.cypher(
                    q, params=(limit,), graph=self._graph
                )
                for record in result:
                    t = self._record_to_triplet(record)
                    if t is not None and t[1].label not in ignore_set:
                        triplets.append(t)

        except IntegrationError:
            raise
        except Exception as exc:
            raise IntegrationError(f"get_rel_map failed: {exc}") from exc

        return triplets

    # ------------------------------------------------------------------
    # Delete
    # ------------------------------------------------------------------

    def delete(
        self,
        entity_names: list[str] | None = None,
        relation_names: list[str] | None = None,
        properties: dict[str, Any] | None = None,
        ids: list[str] | None = None,
    ) -> None:
        """Delete nodes/edges from AGE and corresponding rows from embedding table."""
        try:
            if ids:
                for nid in ids:
                    # Delete from AGE
                    q = "MATCH (n {id: %s}) DETACH DELETE n"
                    self._driver.cypher(q, params=(nid,), graph=self._graph)
                    # Delete from embedding table
                    self._driver.sql(
                        f"DELETE FROM {self._embedding_table} WHERE node_id = %s",
                        (nid,),
                    )

            if entity_names:
                for name in entity_names:
                    q = "MATCH (n {name: %s}) DETACH DELETE n"
                    self._driver.cypher(q, params=(name,), graph=self._graph)
                    self._driver.sql(
                        f"DELETE FROM {self._embedding_table} "
                        f"WHERE properties->>'name' = %s",
                        (name,),
                    )

            if relation_names:
                for rel_name in relation_names:
                    _validate_label(rel_name)
                    q = f"MATCH ()-[r:{rel_name}]->() DELETE r"
                    self._driver.cypher(q, graph=self._graph)

            if properties:
                keys = list(properties.keys())
                vals = list(properties.values())
                where_parts = " AND ".join(
                    f"n.{_validate_label(k)} = %s" for k in keys
                )
                q = f"MATCH (n) WHERE {where_parts} DETACH DELETE n"
                self._driver.cypher(q, params=tuple(vals), graph=self._graph)

            self._driver.commit()
        except IntegrationError:
            self._driver.rollback()
            raise
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"delete failed: {exc}") from exc

    # ------------------------------------------------------------------
    # Structured + vector queries
    # ------------------------------------------------------------------

    def structured_query(
        self,
        query: str,
        param_map: dict[str, Any] | None = None,
    ) -> Any:
        """Pass-through to driver.cypher(), return dicts."""
        try:
            params = tuple(param_map.values()) if param_map else None
            result = self._driver.cypher(
                query, params=params, graph=self._graph
            )
            return result.to_dicts()
        except Exception as exc:
            raise IntegrationError(f"structured_query failed: {exc}") from exc

    def vector_query(
        self,
        query: Any,
        **kwargs: Any,
    ) -> tuple[list[Any], list[float]]:
        """Query the companion pgvector table for similar embeddings.

        Args:
            query: A VectorStoreQuery with query_embedding and similarity_top_k.

        Returns:
            Tuple of (list of LabelledNode, list of scores).
        """
        _load_llamaindex_types()
        try:
            embedding = query.query_embedding
            if embedding is None:
                return ([], [])

            k = getattr(query, "similarity_top_k", 10)
            result = self._driver.vector_search(
                self._embedding_table, "embedding", embedding, k=k,
            )

            nodes: list[Any] = []
            scores: list[float] = []
            for record in result:
                node_id = record.get("node_id", "")
                label = record.get("node_label", "Entity")
                text = record.get("text", "")
                distance = record.get("_distance", 0.0)

                import json
                props_raw = record.get("properties", "{}")
                if isinstance(props_raw, str):
                    props = json.loads(props_raw)
                elif isinstance(props_raw, dict):
                    props = props_raw
                else:
                    props = {}

                # Build the appropriate node type
                assert _ChunkNode is not None
                assert _EntityNode is not None
                if text:
                    node = _ChunkNode(
                        id_=str(node_id),
                        label=label,
                        text=text,
                        properties=props,
                    )
                else:
                    name = props.get("name", str(node_id))
                    node = _EntityNode(
                        id_=str(node_id),
                        label=label,
                        name=name,
                        properties=props,
                    )
                nodes.append(node)
                scores.append(float(distance))

            return (nodes, scores)

        except IntegrationError:
            raise
        except Exception as exc:
            raise IntegrationError(f"vector_query failed: {exc}") from exc

    # ------------------------------------------------------------------
    # Schema
    # ------------------------------------------------------------------

    def get_schema(self, refresh: bool = False) -> dict[str, Any]:
        """Return schema dict from driver.get_schema()."""
        if self._schema_cache is None or refresh:
            try:
                self._schema_cache = self._driver.get_schema(graph=self._graph)
            except Exception as exc:
                raise IntegrationError(f"get_schema failed: {exc}") from exc
        return self._schema_cache

    def get_schema_str(self, refresh: bool = False) -> str:
        """Return a formatted string representation of the schema."""
        schema = self.get_schema(refresh=refresh)
        parts: list[str] = []
        for vl in schema.get("vertex_labels", []):
            count = vl.get("estimated_count", "?")
            parts.append(f"Node: {vl['name']} (~{count} rows)")
        for el in schema.get("edge_labels", []):
            count = el.get("estimated_count", "?")
            parts.append(f"Edge: {el['name']} (~{count} rows)")
        return "\n".join(parts)
