"""LangChain integration — VectorStore and GraphStore backed by Piggie.

Implements the LangChain VectorStore and GraphStore interfaces for
building RAG pipelines with AGE graphs and pgvector embeddings.
"""

from __future__ import annotations

import uuid
from typing import Any, ClassVar

from piggie._util import validate_label as _validate_label
from piggie.exceptions import IntegrationError

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

try:
    from langchain_core.documents import Document
    from langchain_core.embeddings import Embeddings
    from langchain_core.vectorstores import VectorStore as LCVectorStore

    _HAS_LANGCHAIN = True
except ImportError:
    _HAS_LANGCHAIN = False

    class Document:  # type: ignore[no-redef]
        def __init__(self, page_content: str = "", metadata: dict[str, Any] | None = None):
            self.page_content = page_content
            self.metadata = metadata or {}

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

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

try:
    from langchain_community.graphs.graph_store import GraphStore as LCGraphStore

    _HAS_GRAPH_STORE = True
except ImportError:
    _HAS_GRAPH_STORE = False

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




# ---------------------------------------------------------------------------
# PiggieVectorStore
# ---------------------------------------------------------------------------


class PiggieVectorStore(LCVectorStore):
    """LangChain-compatible VectorStore backed by pgvector in PostgreSQL.

    Stores documents in a SQL table with text, metadata, and embedding columns.
    Uses the PiggieDriver for vector similarity search.

    Usage::

        from piggie._integrations.langchain import PiggieVectorStore
        vs = PiggieVectorStore(driver, embedding, table="documents")
    """

    _DEFAULT_TABLE: ClassVar[str] = "langchain_documents"
    _DEFAULT_CONTENT_COL: ClassVar[str] = "content"
    _DEFAULT_EMBEDDING_COL: ClassVar[str] = "embedding"
    _DEFAULT_METADATA_COL: ClassVar[str] = "metadata"

    def __init__(
        self,
        driver: Any,
        embedding: Any,
        *,
        table: str = _DEFAULT_TABLE,
        content_column: str = _DEFAULT_CONTENT_COL,
        embedding_column: str = _DEFAULT_EMBEDDING_COL,
        metadata_column: str = _DEFAULT_METADATA_COL,
        create_table: bool = False,
        embedding_dimensions: int | None = None,
    ) -> None:
        self._driver = driver
        self._embedding = embedding
        self._table = table
        self._content_col = content_column
        self._embedding_col = embedding_column
        self._metadata_col = metadata_column

        if create_table:
            self._ensure_table(embedding_dimensions)

    def _ensure_table(self, dimensions: int | None = None) -> None:
        """Create the vector store table if it does not exist."""
        dim_spec = f"({dimensions})" if dimensions else ""
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self._table} ("
            f"id TEXT PRIMARY KEY, "
            f"{self._content_col} TEXT, "
            f"{self._embedding_col} vector{dim_spec}, "
            f"{self._metadata_col} JSONB DEFAULT '{{}}'"
            f")"
        )
        try:
            self._driver.sql(sql)
            self._driver.commit()
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"_ensure_table failed: {exc}") from exc

    @property
    def embeddings(self) -> Any:
        return self._embedding

    def add_texts(
        self,
        texts: list[str],
        metadatas: list[dict[str, Any]] | None = None,
        ids: list[str] | None = None,
        **kwargs: Any,
    ) -> list[str]:
        """Add texts with embeddings to the vector store."""
        import json

        vectors = self._embedding.embed_documents(list(texts))
        result_ids: list[str] = []

        for i, (text, vec) in enumerate(zip(texts, vectors)):
            doc_id = ids[i] if ids else str(uuid.uuid4())
            meta = metadatas[i] if metadatas else {}
            meta_json = json.dumps(meta)
            vec_str = "[" + ",".join(str(v) for v in vec) + "]"

            sql = (
                f"INSERT INTO {self._table} (id, {self._content_col}, "
                f"{self._embedding_col}, {self._metadata_col}) "
                f"VALUES (%s, %s, %s::vector, %s::jsonb)"
            )
            try:
                self._driver.sql(sql, (doc_id, text, vec_str, meta_json))
            except Exception as exc:
                self._driver.rollback()
                raise IntegrationError(f"add_texts failed: {exc}") from exc
            result_ids.append(doc_id)

        self._driver.commit()
        return result_ids

    @classmethod
    def from_texts(
        cls,
        texts: list[str],
        embedding: Any,
        metadatas: list[dict[str, Any]] | None = None,
        *,
        driver: Any = None,
        table: str = _DEFAULT_TABLE,
        **kwargs: Any,
    ) -> PiggieVectorStore:
        """Create a PiggieVectorStore from a list of texts."""
        if driver is None:
            raise IntegrationError("driver is required for from_texts()")
        store = cls(driver, embedding, table=table)
        store.add_texts(texts, metadatas=metadatas)
        return store

    def delete(self, ids: list[str] | None = None, **kwargs: Any) -> bool | None:
        """Delete documents by ID from the vector store."""
        if not ids:
            return False
        placeholders = ", ".join(["%s"] * len(ids))
        sql = f"DELETE FROM {self._table} WHERE id IN ({placeholders})"
        try:
            self._driver.sql(sql, tuple(ids))
            self._driver.commit()
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"delete failed: {exc}") from exc
        return True

    def similarity_search(
        self,
        query: str,
        k: int = 4,
        **kwargs: Any,
    ) -> list[Any]:
        """Search for similar documents by text query."""
        docs_with_scores = self.similarity_search_with_score(query, k=k, **kwargs)
        return [doc for doc, _ in docs_with_scores]

    def similarity_search_with_score(
        self,
        query: str,
        k: int = 4,
        **kwargs: Any,
    ) -> list[tuple[Any, float]]:
        """Search for similar documents, returning (Document, score) tuples."""
        query_vec = self._embedding.embed_query(query)
        return self.similarity_search_by_vector_with_score(query_vec, k=k, **kwargs)

    def similarity_search_by_vector(
        self,
        embedding: list[float],
        k: int = 4,
        **kwargs: Any,
    ) -> list[Any]:
        """Search for similar documents by pre-computed embedding vector."""
        docs_with_scores = self.similarity_search_by_vector_with_score(
            embedding, k=k, **kwargs
        )
        return [doc for doc, _ in docs_with_scores]

    def similarity_search_by_vector_with_score(
        self,
        embedding: list[float],
        k: int = 4,
        **kwargs: Any,
    ) -> list[tuple[Any, float]]:
        """Search by pre-computed vector, returning (Document, score) tuples."""
        import json

        result = self._driver.vector_search(
            self._table, self._embedding_col, embedding, k=k,
        )

        docs: list[tuple[Any, float]] = []
        for record in result:
            content = record.get(self._content_col, "")
            meta_raw = record.get(self._metadata_col, "{}")
            if isinstance(meta_raw, str):
                meta = json.loads(meta_raw)
            elif isinstance(meta_raw, dict):
                meta = meta_raw
            else:
                meta = {}
            distance = record.get("_distance", 0.0)
            doc = Document(page_content=str(content), metadata=meta)
            docs.append((doc, float(distance)))

        return docs


# ---------------------------------------------------------------------------
# PiggieGraphStore (LangChain)
# ---------------------------------------------------------------------------


class PiggieGraphStore(LCGraphStore):
    """LangChain-compatible GraphStore backed by Apache AGE.

    Usage::

        from piggie._integrations.langchain import PiggieGraphStore
        gs = PiggieGraphStore(driver, graph="my_graph")
    """

    def __init__(self, driver: Any, graph: str = "default") -> None:
        self._driver = driver
        self._graph = graph
        self._schema: str = ""
        self._structured_schema: dict[str, Any] = {}

    @property
    def get_schema(self) -> str:
        """Return a string representation of the graph schema."""
        if not self._schema:
            self.refresh_schema()
        return self._schema

    @property
    def get_structured_schema(self) -> dict[str, Any]:
        """Return structured schema with node_props, rel_props, relationships."""
        if not self._structured_schema:
            self.refresh_schema()
        return self._structured_schema

    def refresh_schema(self) -> None:
        """Refresh the cached schema from the database."""
        try:
            schema = self._driver.get_schema(graph=self._graph)
            parts = []
            node_props: dict[str, Any] = {}
            rel_props: dict[str, Any] = {}
            relationships: list[str] = []

            for vl in schema.get("vertex_labels", []):
                name = vl["name"]
                parts.append(f"Node: {name}")
                node_props[name] = list(vl.get("properties", {}).keys()) if vl.get("properties") else []
            for el in schema.get("edge_labels", []):
                name = el["name"]
                parts.append(f"Edge: {name}")
                rel_props[name] = list(el.get("properties", {}).keys()) if el.get("properties") else []
                relationships.append(name)

            self._schema = "\n".join(parts)
            self._structured_schema = {
                "node_props": node_props,
                "rel_props": rel_props,
                "relationships": relationships,
            }
        except Exception as exc:
            raise IntegrationError(f"refresh_schema failed: {exc}") from exc

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

    def add_graph_documents(
        self,
        graph_documents: list[Any],
        include_source: bool = False,
        **kwargs: Any,
    ) -> None:
        """Add graph documents (nodes + relationships) to the graph.

        Expects objects with .nodes (list of Node with id, type, properties)
        and .relationships (list with source, target, type, properties).
        """
        try:
            for doc in graph_documents:
                # Upsert nodes
                for node in getattr(doc, "nodes", []):
                    node_id = getattr(node, "id", None)
                    node_type = _validate_label(getattr(node, "type", "Entity"))
                    props = getattr(node, "properties", {})
                    if node_id:
                        # Build SET clause with parameterized values
                        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:{node_type} {{id: %s}})"
                        params: tuple[Any, ...] = (node_id,)
                        if set_parts:
                            q += f" SET {set_parts}"
                            params = (node_id, *prop_vals)
                        self._driver.cypher(q, params=params, graph=self._graph)

                # Create relationships
                for rel in getattr(doc, "relationships", []):
                    src = getattr(rel, "source", None)
                    tgt = getattr(rel, "target", None)
                    rel_type = _validate_label(getattr(rel, "type", "RELATED"))
                    if src and tgt:
                        src_id = getattr(src, "id", src)
                        tgt_id = getattr(tgt, "id", tgt)
                        src_type = _validate_label(getattr(src, "type", "Entity"))
                        tgt_type = _validate_label(getattr(tgt, "type", "Entity"))
                        q = (
                            f"MATCH (a:{src_type} {{id: %s}}), "
                            f"(b:{tgt_type} {{id: %s}}) "
                            f"MERGE (a)-[:{rel_type}]->(b)"
                        )
                        self._driver.cypher(
                            q, params=(src_id, tgt_id), graph=self._graph
                        )

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