"""
neptune.py
----------
AWS Neptune implementation of GraphBackend (via Gremlin).
"""

from typing import Any, Dict, List, Optional, Tuple

try:
    from gremlin_python.driver import client
except ImportError:
    client = None

from hieromem_core.backends.interface import GraphBackend


class NeptuneBackend(GraphBackend):
    def __init__(self, endpoint: str):
        if not client:
            raise ImportError("gremlinpython is required for NeptuneBackend")
        # Ensure wss:// prefix
        if not endpoint.startswith("wss://"):
            endpoint = f"wss://{endpoint}/gremlin"

        self.client = client.Client(endpoint, "g")

    def add_triples(self, triples: List[Tuple[str, str, str, Dict[str, Any]]]) -> None:
        # Batching gremlin queries is ideal
        for u, rel, v, props in triples:
            # Basic gremlin string construction (unsafe for arbitrary input, standard warning)
            # V().has('id', u).addE(rel).to(V().has('id', v))
            query = (
                f"g.V().has('id', '{u}').fold().coalesce(unfold(), addV('Node').property('id', '{u}'))"
                f".addE('{rel}').to("
                f"  g.V().has('id', '{v}').fold().coalesce(unfold(), addV('Node').property('id', '{v}'))"
                f")"
            )
            for k, val in props.items():
                # Property serialization...
                query += f".property('{k}', '{val}')"

            self.client.submit(query).all().result()

    def search_triples(
        self,
        head: Optional[str] = None,
        relation: Optional[str] = None,
        tail: Optional[str] = None,
    ) -> List[Tuple[str, str, str, Dict[str, Any]]]:
        """
        Search for triples matching the given pattern.

        Args:
            head: Optional head node ID to filter by
            relation: Optional relation/edge label to filter by
            tail: Optional tail node ID to filter by

        Returns:
            List of matching triples as (head, relation, tail, properties)
        """
        # Build Gremlin query based on provided filters
        if head and tail and relation:
            # Specific edge lookup
            query = (
                f"g.V().has('id', '{head}')"
                f".outE('{relation}')"
                f".where(inV().has('id', '{tail}'))"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )
        elif head and relation:
            # Edges from head with specific relation
            query = (
                f"g.V().has('id', '{head}')"
                f".outE('{relation}')"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )
        elif head:
            # All edges from head
            query = (
                f"g.V().has('id', '{head}')"
                f".outE()"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )
        elif tail and relation:
            # Edges to tail with specific relation
            query = (
                f"g.V().has('id', '{tail}')"
                f".inE('{relation}')"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )
        elif tail:
            # All edges to tail
            query = (
                f"g.V().has('id', '{tail}')"
                f".inE()"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )
        elif relation:
            # All edges with specific relation
            query = (
                f"g.E().hasLabel('{relation}')"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )
        else:
            # All edges (use with caution on large graphs)
            query = (
                "g.E().limit(1000)"
                ".project('head', 'rel', 'tail', 'props')"
                ".by(outV().values('id'))"
                ".by(label())"
                ".by(inV().values('id'))"
                ".by(valueMap())"
            )

        result = self.client.submit(query).all().result()

        triples = []
        for row in result:
            head_id = row.get("head", "")
            rel = row.get("rel", "")
            tail_id = row.get("tail", "")
            props = row.get("props", {})
            # Neptune returns properties as lists, flatten single values
            flat_props = {
                k: v[0] if isinstance(v, list) and len(v) == 1 else v
                for k, v in props.items()
            }
            triples.append((head_id, rel, tail_id, flat_props))

        return triples

    def get_neighbors(
        self, node: str, depth: int = 1
    ) -> List[Tuple[str, str, str, Dict[str, Any]]]:
        """
        Get all neighbors of a node up to specified depth.

        Args:
            node: The node ID to get neighbors for
            depth: How many hops to traverse (default 1)

        Returns:
            List of triples connecting to neighbors
        """
        if depth < 1:
            return []

        # Build repeat traversal for multi-hop
        if depth == 1:
            query = (
                f"g.V().has('id', '{node}')"
                f".bothE()"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )
        else:
            # Multi-hop: collect edges along the path
            query = (
                f"g.V().has('id', '{node}')"
                f".repeat(bothE().store('e').otherV()).times({depth})"
                f".cap('e').unfold()"
                f".dedup()"
                f".project('head', 'rel', 'tail', 'props')"
                f".by(outV().values('id'))"
                f".by(label())"
                f".by(inV().values('id'))"
                f".by(valueMap())"
            )

        result = self.client.submit(query).all().result()

        triples = []
        for row in result:
            head_id = row.get("head", "")
            rel = row.get("rel", "")
            tail_id = row.get("tail", "")
            props = row.get("props", {})
            # Neptune returns properties as lists, flatten single values
            flat_props = {
                k: v[0] if isinstance(v, list) and len(v) == 1 else v
                for k, v in props.items()
            }
            triples.append((head_id, rel, tail_id, flat_props))

        return triples

    def delete_triples(self, triples: List[Tuple[str, str, str]]) -> None:
        """Delete specific triples from the graph."""
        for head, rel, tail in triples:
            query = (
                f"g.V().has('id', '{head}')"
                f".outE('{rel}')"
                f".where(inV().has('id', '{tail}'))"
                f".drop()"
            )
            self.client.submit(query).all().result()

    def count(self) -> int:
        """Return the number of edges in the graph."""
        result = self.client.submit("g.E().count()").all().result()
        return result[0] if result else 0

    def clear(self) -> None:
        self.client.submit("g.V().drop()").all().result()
