"""LlamaIndex integration — GraphStore backed by Apache AGE.

Implements the llama_index GraphStore protocol for knowledge-graph RAG.
Triplets are stored as::

    (:Entity {name: subj})-[:PREDICATE {type: rel}]->(:Entity {name: obj})
"""

from __future__ import annotations

from typing import Any

from piggie.exceptions import IntegrationError

try:
    from llama_index.core.graph_stores.types import GraphStore as LIGraphStore

    _HAS_LLAMAINDEX = True
except ImportError:
    _HAS_LLAMAINDEX = False

    class LIGraphStore:  # type: ignore[no-redef]
        """Stub when llama_index is not installed."""


class PiggieGraphStore(LIGraphStore):
    """LlamaIndex-compatible GraphStore backed by Apache AGE + PostgreSQL.

    Usage::

        from piggie._integrations.llamaindex import PiggieGraphStore
        store = PiggieGraphStore(driver, graph="my_graph")
    """

    def __init__(self, driver: Any, graph: str = "default") -> None:
        self._driver = driver
        self._graph = graph
        self._schema: str = ""

    @property
    def client(self) -> Any:
        """Return the underlying PiggieDriver."""
        return self._driver

    def get(
        self,
        subj: str,
        depth: int = 2,
        limit: int = 30,
    ) -> list[list[str]]:
        """Get triplets for a subject up to a given depth.

        Returns:
            List of [subject, predicate, object] triplets.
        """
        triplets: list[list[str]] = []
        try:
            query = (
                f"MATCH (s:Entity {{name: %s}})-[r:PREDICATE*1..{depth}]->(o:Entity) "
                f"RETURN s.name, r, o.name LIMIT {limit}"
            )
            result = self._driver.cypher(query, params=(subj,), graph=self._graph)
            for record in result:
                s_name = record[0]
                o_name = record[2]
                # r is a list of edges for VLE; extract each hop
                r_val = record[1]
                if isinstance(r_val, list):
                    for edge in r_val:
                        rel = edge.properties.get("type", "RELATED")
                        triplets.append([str(s_name), str(rel), str(o_name)])
                else:
                    rel = r_val.properties.get("type", "RELATED") if hasattr(r_val, "properties") else "RELATED"
                    triplets.append([str(s_name), str(rel), str(o_name)])
        except Exception as exc:
            raise IntegrationError(f"LlamaIndex get() failed: {exc}") from exc
        return triplets

    def get_rel_map(
        self,
        subjs: list[str] | None = None,
        depth: int = 2,
        limit: int = 30,
    ) -> dict[str, list[list[str]]]:
        """Get a relationship map for given subjects.

        Returns:
            Dict mapping subject names to lists of [predicate, object] pairs.
        """
        rel_map: dict[str, list[list[str]]] = {}

        if subjs is None:
            try:
                result = self._driver.cypher(
                    "MATCH (s:Entity) RETURN DISTINCT s.name", graph=self._graph
                )
                subjs = [str(r[0]) for r in result]
            except Exception as exc:
                raise IntegrationError(f"Failed to list subjects: {exc}") from exc

        for subj in subjs:
            triplets = self.get(subj, depth=depth, limit=limit)
            if triplets:
                rel_map[subj] = [[t[1], t[2]] for t in triplets]

        return rel_map

    def upsert_triplet(self, subj: str, rel: str, obj: str) -> None:
        """Upsert a single triplet using MERGE."""
        try:
            query = (
                "MERGE (s:Entity {name: %s}) "
                "MERGE (o:Entity {name: %s}) "
                "MERGE (s)-[:PREDICATE {type: %s}]->(o)"
            )
            self._driver.cypher(query, params=(subj, obj, rel), graph=self._graph)
            self._driver.commit()
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"upsert_triplet failed: {exc}") from exc

    def delete(self, subj: str, rel: str, obj: str) -> None:
        """Delete a specific triplet."""
        try:
            query = (
                "MATCH (s:Entity {name: %s})-[r:PREDICATE {type: %s}]->(o:Entity {name: %s}) "
                "DELETE r"
            )
            self._driver.cypher(query, params=(subj, rel, obj), graph=self._graph)
            self._driver.commit()
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"delete triplet failed: {exc}") from exc

    def persist(
        self,
        persist_path: str = "",
        fs: Any = None,
    ) -> None:
        """No-op — data is already persisted in PostgreSQL."""

    def get_schema(self, refresh: bool = False) -> str:
        """Return a string representation of the graph schema."""
        if not self._schema or refresh:
            try:
                schema = self._driver.get_schema(graph=self._graph)
                parts = []
                for vl in schema.get("vertex_labels", []):
                    parts.append(f"Node: {vl['name']} (~{vl['estimated_count']} rows)")
                for el in schema.get("edge_labels", []):
                    parts.append(f"Edge: {el['name']} (~{el['estimated_count']} rows)")
                self._schema = "\n".join(parts)
            except Exception as exc:
                raise IntegrationError(f"get_schema failed: {exc}") from exc
        return self._schema

    def query(self, query: str, param_map: dict[str, Any] | None = None) -> str:
        """Execute a Cypher query and return results as a string."""
        try:
            params = tuple(param_map.values()) if param_map else None
            result = self._driver.cypher(query, params=params, graph=self._graph)
            return str(result.to_dicts())
        except Exception as exc:
            raise IntegrationError(f"query failed: {exc}") from exc
