"""PyTorch Geometric integration — FeatureStore and GraphStore backed by AGE + PostgreSQL.

Feature tensors are stored in a dedicated PostgreSQL table (``{graph}_node_features``)
with bytea serialization. Edge topology (COO format) is extracted from AGE.

Node IDs are remapped from AGE's 64-bit IDs to contiguous 0-based indices as
required by PyG.
"""

from __future__ import annotations

import pickle
from typing import Any

from piggie.exceptions import IntegrationError

# ---------------------------------------------------------------------------
# Conditional imports
# ---------------------------------------------------------------------------

try:
    import torch as _torch  # noqa: F401 — needed for torch_geometric imports
    from torch_geometric.data.feature_store import FeatureStore, TensorAttr
    from torch_geometric.data.graph_store import EdgeAttr
    from torch_geometric.data.graph_store import GraphStore as PyGGraphStore

    _HAS_PYG = True
except ImportError:
    _HAS_PYG = False

    # Stubs for type checking when PyG is not installed
    class TensorAttr:  # type: ignore[no-redef]
        def __init__(self, group_name: str = "", attr_name: str = "", index: Any = None):
            self.group_name = group_name
            self.attr_name = attr_name
            self.index = index

    class EdgeAttr:  # type: ignore[no-redef]
        def __init__(self, edge_type: Any = None, layout: str = "coo", is_sorted: bool = False, size: Any = None):
            self.edge_type = edge_type
            self.layout = layout
            self.is_sorted = is_sorted
            self.size = size

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

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


def _ensure_feature_table(driver: Any, graph: str) -> str:
    """Create the feature storage table if it doesn't exist. Returns table name."""
    table = f"{graph}_node_features"
    driver.sql(
        f"CREATE TABLE IF NOT EXISTS {table} ("
        f"  node_id BIGINT NOT NULL,"
        f"  node_type TEXT NOT NULL,"
        f"  feature_name TEXT NOT NULL,"
        f"  tensor BYTEA NOT NULL,"
        f"  PRIMARY KEY (node_id, node_type, feature_name)"
        f")"
    )
    driver.commit()
    return table


# ---------------------------------------------------------------------------
# PiggieFeatureStore
# ---------------------------------------------------------------------------


class PiggieFeatureStore(FeatureStore):
    """PyG FeatureStore backed by a PostgreSQL table.

    Feature tensors are serialized via pickle and stored as bytea.

    Usage::

        from piggie._integrations.pyg import PiggieFeatureStore
        fs = PiggieFeatureStore(driver, graph="my_graph")
        fs._put_tensor(tensor, TensorAttr("Person", "x"))
        t = fs._get_tensor(TensorAttr("Person", "x"))
    """

    def __init__(self, driver: Any, graph: str = "default") -> None:
        if _HAS_PYG:
            super().__init__()
        self._driver = driver
        self._graph = graph
        self._table = _ensure_feature_table(driver, graph)

    def _put_tensor(self, tensor: Any, attr: TensorAttr) -> bool:
        """Store a feature tensor."""
        try:
            data = pickle.dumps(tensor.cpu().numpy())
        except Exception as exc:
            raise IntegrationError(f"Failed to serialize tensor: {exc}") from exc

        try:
            self._driver.sql(
                f"INSERT INTO {self._table} (node_id, node_type, feature_name, tensor) "
                f"VALUES (%s, %s, %s, %s) "
                f"ON CONFLICT (node_id, node_type, feature_name) "
                f"DO UPDATE SET tensor = EXCLUDED.tensor",
                (
                    attr.index if attr.index is not None else 0,
                    attr.group_name,
                    attr.attr_name,
                    data,
                ),
            )
            self._driver.commit()
            return True
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"_put_tensor failed: {exc}") from exc

    def _get_tensor(self, attr: TensorAttr) -> Any:
        """Retrieve a feature tensor."""
        try:
            import torch

            if attr.index is not None:
                result = self._driver.sql(
                    f"SELECT tensor FROM {self._table} "
                    f"WHERE node_id = %s AND node_type = %s AND feature_name = %s",
                    (attr.index, attr.group_name, attr.attr_name),
                )
            else:
                result = self._driver.sql(
                    f"SELECT tensor FROM {self._table} "
                    f"WHERE node_type = %s AND feature_name = %s "
                    f"ORDER BY node_id",
                    (attr.group_name, attr.attr_name),
                )

            if not result:
                return None

            if attr.index is not None:
                data = result.single()[0]
                if isinstance(data, memoryview):
                    data = bytes(data)
                return torch.from_numpy(pickle.loads(data))  # noqa: S301
            else:
                tensors = []
                for record in result:
                    data = record[0]
                    if isinstance(data, memoryview):
                        data = bytes(data)
                    tensors.append(torch.from_numpy(pickle.loads(data)))  # noqa: S301
                if tensors:
                    return torch.stack(tensors)
                return None
        except IntegrationError:
            raise
        except Exception as exc:
            raise IntegrationError(f"_get_tensor failed: {exc}") from exc

    def _remove_tensor(self, attr: TensorAttr) -> bool:
        """Remove a stored tensor."""
        try:
            if attr.index is not None:
                self._driver.sql(
                    f"DELETE FROM {self._table} "
                    f"WHERE node_id = %s AND node_type = %s AND feature_name = %s",
                    (attr.index, attr.group_name, attr.attr_name),
                )
            else:
                self._driver.sql(
                    f"DELETE FROM {self._table} "
                    f"WHERE node_type = %s AND feature_name = %s",
                    (attr.group_name, attr.attr_name),
                )
            self._driver.commit()
            return True
        except Exception as exc:
            self._driver.rollback()
            raise IntegrationError(f"_remove_tensor failed: {exc}") from exc

    def get_all_tensor_attrs(self) -> list[TensorAttr]:
        """List all stored tensor attributes."""
        result = self._driver.sql(
            f"SELECT DISTINCT node_type, feature_name FROM {self._table} "
            f"ORDER BY node_type, feature_name"
        )
        return [
            TensorAttr(group_name=r[0], attr_name=r[1])
            for r in result
        ]


# ---------------------------------------------------------------------------
# PiggieGraphStore (PyG)
# ---------------------------------------------------------------------------


class PiggieGraphStore(PyGGraphStore):
    """PyG GraphStore backed by AGE edge topology.

    Extracts edge indices in COO format from AGE graphs. Node IDs are
    remapped from AGE's large 64-bit IDs to contiguous 0-based indices.

    Usage::

        from piggie._integrations.pyg import PiggieGraphStore
        gs = PiggieGraphStore(driver, graph="my_graph")
        edge_index = gs._get_edge_index(EdgeAttr(("Person", "KNOWS", "Person")))
    """

    def __init__(self, driver: Any, graph: str = "default") -> None:
        if _HAS_PYG:
            super().__init__()
        self._driver = driver
        self._graph = graph
        self._node_id_maps: dict[str, dict[int, int]] = {}

    def _build_node_id_map(self, label: str) -> dict[int, int]:
        """Build a mapping from AGE vertex IDs to contiguous 0-based indices."""
        if label in self._node_id_maps:
            return self._node_id_maps[label]

        result = self._driver.cypher(
            f"MATCH (n:{label}) RETURN id(n) AS nid ORDER BY id(n)",
            graph=self._graph,
        )
        id_map = {int(record["nid"]): idx for idx, record in enumerate(result)}
        self._node_id_maps[label] = id_map
        return id_map

    def _put_edge_index(self, edge_index: Any, edge_attr: EdgeAttr) -> bool:
        """Store edge index — no-op since topology is managed by AGE."""
        return True

    def _get_edge_index(self, edge_attr: EdgeAttr) -> Any:
        """Extract edge index in COO format [2, num_edges] from AGE.

        Args:
            edge_attr: EdgeAttr with edge_type = (src_label, rel_type, dst_label).

        Returns:
            torch.LongTensor of shape [2, num_edges] or None.
        """
        try:
            import torch
        except ImportError:
            raise ImportError(
                "torch is required. Install with: pip install piggie[pyg]"
            ) from None

        edge_type = edge_attr.edge_type
        if edge_type is None or len(edge_type) != 3:
            raise IntegrationError(
                "edge_type must be a (src_label, rel_type, dst_label) tuple"
            )

        src_label, rel_type, dst_label = edge_type
        src_map = self._build_node_id_map(src_label)
        dst_map = self._build_node_id_map(dst_label)

        result = self._driver.cypher(
            f"MATCH (a:{src_label})-[r:{rel_type}]->(b:{dst_label}) "
            f"RETURN id(a) AS src, id(b) AS dst",
            graph=self._graph,
        )

        src_ids = []
        dst_ids = []
        for record in result:
            src_age_id = int(record["src"])
            dst_age_id = int(record["dst"])
            if src_age_id in src_map and dst_age_id in dst_map:
                src_ids.append(src_map[src_age_id])
                dst_ids.append(dst_map[dst_age_id])

        if not src_ids:
            return None

        return torch.LongTensor([src_ids, dst_ids])

    def _remove_edge_index(self, edge_attr: EdgeAttr) -> bool:
        """Remove edge index — no-op since topology is managed by AGE."""
        return True

    def get_all_edge_attrs(self) -> list[EdgeAttr]:
        """List all edge types in the graph."""
        schema = self._driver.get_schema(graph=self._graph)
        vertex_labels = [lbl["name"] for lbl in schema.get("vertex_labels", [])]
        edge_labels = [lbl["name"] for lbl in schema.get("edge_labels", [])]

        attrs = []
        for rel_type in edge_labels:
            # For each edge type, determine connected node types
            try:
                result = self._driver.cypher(
                    f"MATCH (a)-[r:{rel_type}]->(b) "
                    f"RETURN DISTINCT labels(a) AS src_label, labels(b) AS dst_label "
                    f"LIMIT 1",
                    graph=self._graph,
                    columns=["src_label", "dst_label"],
                )
                if result:
                    # AGE labels() returns a list; take first element
                    record = result.records[0]
                    src_lbl = record["src_label"]
                    dst_lbl = record["dst_label"]
                    if isinstance(src_lbl, list):
                        src_lbl = src_lbl[0] if src_lbl else ""
                    if isinstance(dst_lbl, list):
                        dst_lbl = dst_lbl[0] if dst_lbl else ""
                    attrs.append(
                        EdgeAttr(
                            edge_type=(str(src_lbl), rel_type, str(dst_lbl)),
                            layout="coo",
                        )
                    )
            except Exception:
                # Fallback: pair with all vertex labels
                for src in vertex_labels:
                    for dst in vertex_labels:
                        attrs.append(
                            EdgeAttr(edge_type=(src, rel_type, dst), layout="coo")
                        )
                break

        return attrs
