"""PiggieDriver — the main entry point for graph + vector operations."""

from __future__ import annotations

from typing import Any

import psycopg
from psycopg.client_cursor import ClientCursor

from piggie._cypher.builder import CypherBuilder
from piggie._cypher.wrapper import execute_cypher
from piggie.connection import connect, drop_graph, ensure_graph, list_graphs
from piggie.exceptions import CypherError, PiggieError
from piggie.result import Result


class _Transaction:
    """Context manager that commits on success, rolls back on exception."""

    def __init__(self, driver: PiggieDriver) -> None:
        self._driver = driver

    def __enter__(self) -> PiggieDriver:
        return self._driver

    def __exit__(self, exc_type: type | None, exc_val: BaseException | None, exc_tb: Any) -> None:
        if exc_type is None:
            self._driver.commit()
        else:
            self._driver.rollback()


class PiggieDriver:
    """High-level driver for Apache AGE + pgvector on PostgreSQL.

    Usage::

        import piggie
        db = piggie.connect("postgresql://localhost:5433/piggie", graph="my_graph")
        result = db.cypher("MATCH (n:Person) RETURN n.name, n.age")
        df = result.to_df()
        db.close()

    Or as a context manager::

        with piggie.connect(dsn, graph="g") as db:
            df = db.cypher("MATCH (n) RETURN n").to_df()
    """

    def __init__(self, dsn: str = "", *, graph: str | None = None, **kwargs: Any) -> None:
        self._conn = connect(dsn, graph=graph, **kwargs)
        self._graph = graph

    @property
    def connection(self) -> psycopg.Connection[Any]:
        """The underlying psycopg connection."""
        return self._conn

    @property
    def graph(self) -> str | None:
        """The default graph name."""
        return self._graph

    def close(self) -> None:
        """Close the database connection."""
        if self._conn and not self._conn.closed:
            self._conn.close()

    def __enter__(self) -> PiggieDriver:
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

    def cypher(
        self,
        query: str,
        *,
        params: tuple[Any, ...] | None = None,
        columns: list[str] | None = None,
        graph: str | None = None,
    ) -> Result:
        """Execute a Cypher query and return a Result.

        Args:
            query: Cypher query string.
            params: Parameters for %s placeholder substitution.
            columns: Explicit column names. Auto-extracted from RETURN if None.
            graph: Graph name override. Uses the default graph if None.

        Returns:
            A Result object supporting .to_df(), .to_dicts(), etc.
        """
        graph_name = graph or self._graph
        if not graph_name:
            raise PiggieError("No graph specified. Pass graph= to connect() or cypher().")

        try:
            rows, cols = execute_cypher(
                self._conn, graph_name, query, params=params, columns=columns
            )
        except PiggieError:
            raise
        except Exception as exc:
            self._conn.rollback()
            raise CypherError(f"Cypher execution failed: {exc}") from exc

        return Result(rows, cols)

    def sql(self, query: str, params: tuple[Any, ...] | None = None) -> Result:
        """Execute a raw SQL query and return a Result.

        Args:
            query: SQL query string.
            params: Query parameters.

        Returns:
            A Result object.
        """
        cursor = self._conn.cursor()
        try:
            cursor.execute(query, params)
            if cursor.description:
                columns = [desc[0] for desc in cursor.description]
                rows = cursor.fetchall()
            else:
                columns = []
                rows = []
        except Exception as exc:
            self._conn.rollback()
            raise PiggieError(f"SQL execution failed: {exc}") from exc

        return Result(rows, columns)

    def commit(self) -> None:
        """Commit the current transaction."""
        self._conn.commit()

    def rollback(self) -> None:
        """Roll back the current transaction."""
        self._conn.rollback()

    def transaction(self) -> _Transaction:
        """Context manager for an explicit transaction block.

        Usage::

            with db.transaction():
                db.cypher("CREATE (:Person {name: 'Alice'})")
                db.cypher("CREATE (:Person {name: 'Bob'})")
            # auto-committed on success, rolled back on exception
        """
        return _Transaction(self)

    def set_graph(self, graph_name: str) -> None:
        """Switch to a different graph, creating it if needed."""
        ensure_graph(self._conn, graph_name)
        self._graph = graph_name

    def list_graphs(self) -> list[str]:
        """List all AGE graphs in the database."""
        return list_graphs(self._conn)

    def drop_graph(self, graph_name: str) -> None:
        """Drop a graph and all its data."""
        drop_graph(self._conn, graph_name)
        if self._graph == graph_name:
            self._graph = None

    def vector_search(
        self,
        table: str,
        column: str,
        query_vector: list[float],
        *,
        k: int = 10,
        metric: str = "cosine",
        where: str | None = None,
        columns: list[str] | None = None,
    ) -> Result:
        """Search for similar vectors using pgvector.

        Args:
            table: Table containing the vector column.
            column: Name of the vector column.
            query_vector: Query vector as a list of floats.
            k: Number of nearest neighbors to return.
            metric: Distance metric: "cosine", "l2", or "inner_product".
            where: Optional SQL WHERE clause (without the WHERE keyword).
            columns: Columns to return. Defaults to all (*).

        Returns:
            Result with rows ordered by similarity.
        """
        from piggie._vector.search import vector_search
        return vector_search(
            self._conn, table, column, query_vector,
            k=k, metric=metric, where=where, columns=columns,
        )

    def create_vector_index(
        self,
        table: str,
        column: str,
        *,
        method: str = "hnsw",
        metric: str = "cosine",
        name: str | None = None,
        **kwargs: Any,
    ) -> None:
        """Create a vector index on a table column.

        Args:
            table: Table name.
            column: Vector column name.
            method: Index method: "hnsw" or "ivfflat".
            metric: Distance metric: "cosine", "l2", or "inner_product".
            name: Optional index name.
            **kwargs: Additional index parameters (m, ef_construction, lists).
        """
        from piggie._vector.index import create_vector_index
        create_vector_index(
            self._conn, table, column,
            method=method, metric=metric, name=name, **kwargs,
        )

    def hybrid_search(
        self,
        *,
        cypher: str | None = None,
        cypher_params: tuple[Any, ...] | None = None,
        vector_table: str,
        vector_column: str,
        query_vector: list[float],
        k: int = 10,
        metric: str = "cosine",
        join_column: str = "id",
        graph: str | None = None,
        mode: str = "graph_then_vector",
    ) -> Result:
        """Execute a hybrid graph + vector search.

        Combines Cypher graph traversal with pgvector similarity search.

        Args:
            cypher: Cypher query that produces candidate rows.
            cypher_params: Parameters for the Cypher query.
            vector_table: Table with vector embeddings.
            vector_column: Vector column name.
            query_vector: Query vector.
            k: Number of results.
            metric: Distance metric.
            join_column: Column to join graph results with vector table.
            graph: Graph name override.
            mode: "graph_then_vector" or "vector_then_graph".

        Returns:
            Result combining graph and vector data.
        """
        from piggie._hybrid.search import hybrid_search
        return hybrid_search(
            driver=self,
            cypher=cypher,
            cypher_params=cypher_params,
            vector_table=vector_table,
            vector_column=vector_column,
            query_vector=query_vector,
            k=k,
            metric=metric,
            join_column=join_column,
            graph=graph,
            mode=mode,
        )

    def get_schema(self, graph: str | None = None) -> dict[str, Any]:
        """Get schema information for a graph.

        Returns vertex labels, edge labels, and vector indexes.

        Args:
            graph: Graph name. Uses default if None.

        Returns:
            Dict with keys: vertex_labels, edge_labels, vector_indexes.
        """
        from piggie._schema.introspect import get_schema
        graph_name = graph or self._graph
        if not graph_name:
            raise PiggieError("No graph specified.")
        return get_schema(self._conn, graph_name)

    # ------------------------------------------------------------------
    # Query builder support
    # ------------------------------------------------------------------

    def execute(self, builder: CypherBuilder, *, graph: str | None = None) -> Result:
        """Execute a query builder and return a Result.

        Args:
            builder: A Match, Create, or Merge builder instance.
            graph: Graph name override.

        Returns:
            A Result object.
        """
        cypher, params = builder.build()
        return self.cypher(cypher, params=params or None, graph=graph)

    # ------------------------------------------------------------------
    # Bulk operations
    # ------------------------------------------------------------------

    def load_nodes(
        self,
        df: Any,
        label: str,
        *,
        graph: str | None = None,
        id_col: str | None = None,
        property_cols: list[str] | None = None,
        batch_size: int = 1000,
    ) -> int:
        """Bulk-load nodes from a DataFrame.

        See :func:`piggie._bulk.ingest.load_nodes` for full docs.
        """
        from piggie._bulk.ingest import load_nodes
        return load_nodes(
            self, df, graph or self._graph or "",
            label, id_col=id_col, property_cols=property_cols,
            batch_size=batch_size,
        )

    def load_edges(
        self,
        df: Any,
        rel_type: str,
        source_col: str,
        source_label: str,
        target_col: str,
        target_label: str,
        *,
        source_key: str | None = None,
        target_key: str | None = None,
        graph: str | None = None,
        property_cols: list[str] | None = None,
        batch_size: int = 1000,
    ) -> int:
        """Bulk-load edges from a DataFrame.

        See :func:`piggie._bulk.ingest.load_edges` for full docs.
        """
        from piggie._bulk.ingest import load_edges
        return load_edges(
            self, df, graph or self._graph or "",
            rel_type, source_col, source_label, target_col, target_label,
            source_key=source_key, target_key=target_key,
            property_cols=property_cols, batch_size=batch_size,
        )

    def export_nodes(
        self,
        label: str,
        *,
        graph: str | None = None,
        properties: list[str] | None = None,
        limit: int | None = None,
    ) -> Any:
        """Export nodes to a DataFrame.

        See :func:`piggie._bulk.export.export_nodes` for full docs.
        """
        from piggie._bulk.export import export_nodes
        return export_nodes(
            self, graph or self._graph or "",
            label, properties=properties, limit=limit,
        )

    def export_edges(
        self,
        rel_type: str,
        *,
        graph: str | None = None,
        properties: list[str] | None = None,
        limit: int | None = None,
    ) -> Any:
        """Export edges to a DataFrame.

        See :func:`piggie._bulk.export.export_edges` for full docs.
        """
        from piggie._bulk.export import export_edges
        return export_edges(
            self, graph or self._graph or "",
            rel_type, properties=properties, limit=limit,
        )

    # ------------------------------------------------------------------
    # NetworkX / Graph algorithms
    # ------------------------------------------------------------------

    def to_networkx(
        self,
        *,
        graph: str | None = None,
        labels: list[str] | None = None,
        rel_types: list[str] | None = None,
        directed: bool = True,
        node_properties: bool = True,
        edge_properties: bool = True,
    ) -> Any:
        """Convert the AGE graph to a NetworkX graph object.

        See :func:`piggie._integrations.networkx.to_networkx` for full docs.
        """
        from piggie._integrations.networkx import to_networkx
        return to_networkx(
            self, graph or self._graph or "",
            labels=labels, rel_types=rel_types, directed=directed,
            node_properties=node_properties, edge_properties=edge_properties,
        )

    def pagerank(
        self,
        *,
        graph: str | None = None,
        label: str | None = None,
        rel_type: str | None = None,
        damping: float = 0.85,
    ) -> Any:
        """Compute PageRank over the graph. Returns a DataFrame."""
        from piggie._algorithms.ranking import pagerank
        return pagerank(
            self, graph or self._graph or "",
            label=label, rel_type=rel_type, damping=damping,
        )

    def communities(
        self,
        *,
        graph: str | None = None,
        method: str = "louvain",
        label: str | None = None,
        rel_type: str | None = None,
    ) -> Any:
        """Detect communities in the graph. Returns a DataFrame."""
        from piggie._algorithms.community import community_detection
        return community_detection(
            self, graph or self._graph or "",
            method=method, label=label, rel_type=rel_type,
        )

    def shortest_path(
        self,
        source: int,
        target: int,
        *,
        graph: str | None = None,
        weight: str | None = None,
        label: str | None = None,
        rel_type: str | None = None,
    ) -> list[int]:
        """Find shortest path between two nodes. Returns list of node IDs."""
        from piggie._algorithms.paths import shortest_path
        return shortest_path(
            self, graph or self._graph or "",
            source, target,
            weight=weight, label=label, rel_type=rel_type,
        )

    def centrality(
        self,
        *,
        graph: str | None = None,
        measure: str = "degree",
        label: str | None = None,
        rel_type: str | None = None,
    ) -> Any:
        """Compute centrality measures. Returns a DataFrame."""
        from piggie._algorithms.centrality import centrality
        return centrality(
            self, graph or self._graph or "",
            measure=measure, label=label, rel_type=rel_type,
        )

    def all_shortest_paths(
        self,
        source: int,
        target: int,
        *,
        graph: str | None = None,
        weight: str | None = None,
        label: str | None = None,
        rel_type: str | None = None,
    ) -> list[list[int]]:
        """Find all shortest paths between two nodes."""
        from piggie._algorithms.paths import all_shortest_paths
        return all_shortest_paths(
            self, graph or self._graph or "",
            source, target,
            weight=weight, label=label, rel_type=rel_type,
        )

    def similarity(
        self,
        *,
        measure: str = "jaccard",
        graph: str | None = None,
        node_pairs: list[tuple[int, int]] | None = None,
        label: str | None = None,
        rel_type: str | None = None,
    ) -> Any:
        """Compute similarity scores between node pairs. Returns a DataFrame."""
        from piggie._algorithms import similarity as sim_mod
        graph_name = graph or self._graph or ""
        if measure == "jaccard":
            return sim_mod.jaccard_similarity(
                self, graph_name, node_pairs, label=label, rel_type=rel_type,
            )
        elif measure == "adamic_adar":
            return sim_mod.adamic_adar(
                self, graph_name, node_pairs, label=label, rel_type=rel_type,
            )
        else:
            from piggie.exceptions import AlgorithmError
            raise AlgorithmError(
                f"Unknown similarity measure: {measure!r}. "
                "Valid measures: jaccard, adamic_adar"
            )

    def link_prediction(
        self,
        *,
        measure: str = "preferential_attachment",
        graph: str | None = None,
        node_pairs: list[tuple[int, int]] | None = None,
        label: str | None = None,
        rel_type: str | None = None,
    ) -> Any:
        """Compute link prediction scores. Returns a DataFrame."""
        from piggie._algorithms import link_prediction as lp_mod
        graph_name = graph or self._graph or ""
        if measure == "preferential_attachment":
            return lp_mod.preferential_attachment(
                self, graph_name, node_pairs, label=label, rel_type=rel_type,
            )
        elif measure == "resource_allocation":
            return lp_mod.resource_allocation(
                self, graph_name, node_pairs, label=label, rel_type=rel_type,
            )
        else:
            from piggie.exceptions import AlgorithmError
            raise AlgorithmError(
                f"Unknown link prediction measure: {measure!r}. "
                "Valid measures: preferential_attachment, resource_allocation"
            )

    def node2vec(
        self,
        *,
        graph: str | None = None,
        dimensions: int = 64,
        walk_length: int = 30,
        num_walks: int = 200,
        p: float = 1.0,
        q: float = 1.0,
        workers: int = 1,
        label: str | None = None,
        rel_type: str | None = None,
    ) -> Any:
        """Compute node2vec embeddings. Returns a DataFrame."""
        from piggie._algorithms.embeddings import node2vec_embed
        return node2vec_embed(
            self, graph or self._graph or "",
            dimensions=dimensions, walk_length=walk_length,
            num_walks=num_walks, p=p, q=q, workers=workers,
            label=label, rel_type=rel_type,
        )

    def __repr__(self) -> str:
        status = "closed" if self._conn.closed else "connected"
        return f"PiggieDriver(graph={self._graph!r}, status={status!r})"


class PiggiePool:
    """Connection pool for multi-threaded / web applications.

    Wraps psycopg_pool.ConnectionPool with AGE setup on each connection.

    Usage::

        pool = piggie.connect_pool("postgresql://...", graph="g", min_size=2, max_size=10)

        # Get a driver from the pool
        with pool.connection() as db:
            result = db.cypher("MATCH (n) RETURN n LIMIT 10")

        pool.close()
    """

    def __init__(
        self,
        dsn: str = "",
        *,
        graph: str | None = None,
        min_size: int = 2,
        max_size: int = 10,
        **kwargs: Any,
    ) -> None:
        try:
            from psycopg_pool import ConnectionPool
        except ImportError:
            raise ImportError(
                "psycopg_pool is required for connection pooling. "
                "Install with: pip install psycopg_pool"
            ) from None

        self._graph = graph

        def configure(conn: psycopg.Connection[Any]) -> None:
            from piggie.connection import _setup_age
            _setup_age(conn)
            conn.commit()  # pool requires clean (non-INTRANS) state after configure

        self._pool = ConnectionPool(
            dsn,
            min_size=min_size,
            max_size=max_size,
            open=True,
            configure=configure,
            kwargs={"cursor_factory": ClientCursor, **kwargs},
        )

    def connection(self) -> _PooledConnection:
        """Get a PiggieDriver backed by a pooled connection.

        Usage::

            with pool.connection() as db:
                db.cypher("MATCH (n) RETURN n")
        """
        return _PooledConnection(self._pool, self._graph)

    def close(self) -> None:
        """Close the pool and all connections."""
        self._pool.close()

    def __repr__(self) -> str:
        return f"PiggiePool(graph={self._graph!r})"


class _PooledConnection:
    """Context manager that checks out a connection from the pool."""

    def __init__(self, pool: Any, graph: str | None) -> None:
        self._pool = pool
        self._graph = graph
        self._ctx: Any = None
        self._driver: PiggieDriver | None = None

    def __enter__(self) -> PiggieDriver:
        self._ctx = self._pool.connection()
        conn = self._ctx.__enter__()
        self._driver = PiggieDriver.__new__(PiggieDriver)
        self._driver._conn = conn
        self._driver._graph = self._graph
        return self._driver

    def __exit__(self, *exc: object) -> None:
        if self._ctx is not None:
            self._ctx.__exit__(*exc)
