"""
ALMA Graph Memory - Memgraph Backend.

Memgraph implementation of the GraphBackend interface.
Memgraph is compatible with Neo4j's Bolt protocol, so the neo4j Python driver works with it.
"""

import json
import logging
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from alma.graph.base import GraphBackend
from alma.graph.store import Entity, Relationship

logger = logging.getLogger(__name__)


class MemgraphBackend(GraphBackend):
    """
    Memgraph graph database backend.

    Memgraph is an in-memory graph database compatible with Neo4j's Bolt protocol.
    This backend uses the neo4j Python driver.

    Requires neo4j Python driver: pip install neo4j

    Example usage:
        backend = MemgraphBackend(
            uri="bolt://localhost:7687",
            username="",
            password=""
        )
        backend.add_entity(entity)
        backend.close()

    Note: Memgraph typically doesn't require authentication by default,
    but it can be configured. Use empty strings for username/password if
    authentication is disabled.
    """

    def __init__(
        self,
        uri: str = "bolt://localhost:7687",
        username: str = "",
        password: str = "",
        database: str = "memgraph",
    ):
        """
        Initialize Memgraph connection.

        Args:
            uri: Memgraph connection URI (bolt://)
            username: Database username (empty if auth disabled)
            password: Database password (empty if auth disabled)
            database: Database name (default: "memgraph")
        """
        self.uri = uri
        self.username = username
        self.password = password
        self.database = database
        self._driver = None

    def _get_driver(self):
        """Lazy initialization of Memgraph driver (via neo4j package)."""
        if self._driver is None:
            try:
                from neo4j import GraphDatabase

                # Memgraph uses the same Bolt protocol as Neo4j
                if self.username and self.password:
                    self._driver = GraphDatabase.driver(
                        self.uri,
                        auth=(self.username, self.password),
                    )
                else:
                    # No authentication
                    self._driver = GraphDatabase.driver(self.uri)
            except ImportError as err:
                raise ImportError(
                    "neo4j package required for Memgraph graph backend. "
                    "Install with: pip install neo4j"
                ) from err
        return self._driver

    def _run_query(self, query: str, parameters: Optional[Dict] = None) -> List[Dict]:
        """Execute a Cypher query."""
        driver = self._get_driver()
        # Memgraph doesn't use the database parameter in the same way as Neo4j
        # Most Memgraph setups use a single database
        with driver.session() as session:
            result = session.run(query, parameters or {})
            return [dict(record) for record in result]

    def add_entity(self, entity: Entity) -> str:
        """Add or update an entity in Memgraph."""
        # Extract project_id and agent from properties if present
        properties = entity.properties.copy()
        project_id = properties.pop("project_id", None)
        agent = properties.pop("agent", None)

        query = """
        MERGE (e:Entity {id: $id})
        SET e.name = $name,
            e.entity_type = $entity_type,
            e.properties = $properties,
            e.created_at = $created_at
        """
        params = {
            "id": entity.id,
            "name": entity.name,
            "entity_type": entity.entity_type,
            "properties": json.dumps(properties),
            "created_at": entity.created_at.isoformat(),
        }

        # Add optional fields if present
        if project_id:
            query += ", e.project_id = $project_id"
            params["project_id"] = project_id
        if agent:
            query += ", e.agent = $agent"
            params["agent"] = agent

        query += " RETURN e.id as id"

        result = self._run_query(query, params)
        return result[0]["id"] if result else entity.id

    def add_relationship(self, relationship: Relationship) -> str:
        """Add or update a relationship in Memgraph."""
        # Sanitize relationship type for Cypher (remove special characters)
        rel_type = (
            relationship.relation_type.replace("-", "_").replace(" ", "_").upper()
        )

        query = f"""
        MATCH (source:Entity {{id: $source_id}})
        MATCH (target:Entity {{id: $target_id}})
        MERGE (source)-[r:{rel_type}]->(target)
        SET r.id = $id,
            r.properties = $properties,
            r.confidence = $confidence,
            r.created_at = $created_at
        RETURN r.id as id
        """
        result = self._run_query(
            query,
            {
                "id": relationship.id,
                "source_id": relationship.source_id,
                "target_id": relationship.target_id,
                "properties": json.dumps(relationship.properties),
                "confidence": relationship.confidence,
                "created_at": relationship.created_at.isoformat(),
            },
        )
        return result[0]["id"] if result else relationship.id

    def get_entity(self, entity_id: str) -> Optional[Entity]:
        """Get an entity by ID."""
        query = """
        MATCH (e:Entity {id: $id})
        RETURN e.id as id, e.name as name, e.entity_type as entity_type,
               e.properties as properties, e.created_at as created_at,
               e.project_id as project_id, e.agent as agent
        """
        result = self._run_query(query, {"id": entity_id})
        if not result:
            return None

        r = result[0]
        properties = json.loads(r["properties"]) if r["properties"] else {}

        # Add project_id and agent back to properties if present
        if r.get("project_id"):
            properties["project_id"] = r["project_id"]
        if r.get("agent"):
            properties["agent"] = r["agent"]

        return Entity(
            id=r["id"],
            name=r["name"],
            entity_type=r["entity_type"],
            properties=properties,
            created_at=(
                datetime.fromisoformat(r["created_at"])
                if r["created_at"]
                else datetime.now(timezone.utc)
            ),
        )

    def get_entities(
        self,
        entity_type: Optional[str] = None,
        project_id: Optional[str] = None,
        agent: Optional[str] = None,
        limit: int = 100,
    ) -> List[Entity]:
        """Get entities with optional filtering."""
        conditions = []
        params: Dict[str, Any] = {"limit": limit}

        if entity_type:
            conditions.append("e.entity_type = $entity_type")
            params["entity_type"] = entity_type
        if project_id:
            conditions.append("e.project_id = $project_id")
            params["project_id"] = project_id
        if agent:
            conditions.append("e.agent = $agent")
            params["agent"] = agent

        where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""

        query = f"""
        MATCH (e:Entity)
        {where_clause}
        RETURN e.id as id, e.name as name, e.entity_type as entity_type,
               e.properties as properties, e.created_at as created_at,
               e.project_id as project_id, e.agent as agent
        LIMIT $limit
        """

        results = self._run_query(query, params)
        entities = []
        for r in results:
            properties = json.loads(r["properties"]) if r["properties"] else {}
            if r.get("project_id"):
                properties["project_id"] = r["project_id"]
            if r.get("agent"):
                properties["agent"] = r["agent"]

            entities.append(
                Entity(
                    id=r["id"],
                    name=r["name"],
                    entity_type=r["entity_type"],
                    properties=properties,
                    created_at=(
                        datetime.fromisoformat(r["created_at"])
                        if r["created_at"]
                        else datetime.now(timezone.utc)
                    ),
                )
            )
        return entities

    def get_relationships(self, entity_id: str) -> List[Relationship]:
        """Get all relationships for an entity (both directions)."""
        query = """
        MATCH (e:Entity {id: $entity_id})-[r]-(other:Entity)
        RETURN r.id as id,
               CASE WHEN startNode(r).id = $entity_id THEN e.id ELSE other.id END as source_id,
               CASE WHEN endNode(r).id = $entity_id THEN e.id ELSE other.id END as target_id,
               type(r) as relation_type, r.properties as properties,
               r.confidence as confidence, r.created_at as created_at
        """

        results = self._run_query(query, {"entity_id": entity_id})
        relationships = []
        for r in results:
            rel_id = (
                r["id"] or f"{r['source_id']}-{r['relation_type']}-{r['target_id']}"
            )
            relationships.append(
                Relationship(
                    id=rel_id,
                    source_id=r["source_id"],
                    target_id=r["target_id"],
                    relation_type=r["relation_type"],
                    properties=json.loads(r["properties"]) if r["properties"] else {},
                    confidence=r["confidence"] or 1.0,
                    created_at=(
                        datetime.fromisoformat(r["created_at"])
                        if r["created_at"]
                        else datetime.now(timezone.utc)
                    ),
                )
            )
        return relationships

    def search_entities(
        self,
        query: str,
        embedding: Optional[List[float]] = None,
        top_k: int = 10,
    ) -> List[Entity]:
        """
        Search for entities by name.

        Note: Vector similarity search requires Memgraph MAGE with vector operations.
        Falls back to text search if embedding is provided but vector index
        is not available.
        """
        # For now, we do text-based search
        # Vector search can be added when Memgraph MAGE is configured
        cypher = """
        MATCH (e:Entity)
        WHERE toLower(e.name) CONTAINS toLower($query)
        RETURN e.id as id, e.name as name, e.entity_type as entity_type,
               e.properties as properties, e.created_at as created_at,
               e.project_id as project_id, e.agent as agent
        LIMIT $limit
        """

        results = self._run_query(cypher, {"query": query, "limit": top_k})
        entities = []
        for r in results:
            properties = json.loads(r["properties"]) if r["properties"] else {}
            if r.get("project_id"):
                properties["project_id"] = r["project_id"]
            if r.get("agent"):
                properties["agent"] = r["agent"]

            entities.append(
                Entity(
                    id=r["id"],
                    name=r["name"],
                    entity_type=r["entity_type"],
                    properties=properties,
                    created_at=(
                        datetime.fromisoformat(r["created_at"])
                        if r["created_at"]
                        else datetime.now(timezone.utc)
                    ),
                )
            )
        return entities

    def delete_entity(self, entity_id: str) -> bool:
        """Delete an entity and its relationships."""
        query = """
        MATCH (e:Entity {id: $id})
        DETACH DELETE e
        RETURN count(e) as deleted
        """
        result = self._run_query(query, {"id": entity_id})
        return result[0]["deleted"] > 0 if result else False

    def delete_relationship(self, relationship_id: str) -> bool:
        """Delete a specific relationship by ID."""
        query = """
        MATCH ()-[r]-()
        WHERE r.id = $id
        DELETE r
        RETURN count(r) as deleted
        """
        result = self._run_query(query, {"id": relationship_id})
        return result[0]["deleted"] > 0 if result else False

    def close(self) -> None:
        """Close the Memgraph driver connection."""
        if self._driver:
            self._driver.close()
            self._driver = None

    # Additional methods for compatibility with existing GraphStore API

    def find_entities(
        self,
        name: Optional[str] = None,
        entity_type: Optional[str] = None,
        limit: int = 10,
    ) -> List[Entity]:
        """
        Find entities by name or type.

        This method provides compatibility with the existing GraphStore API.
        """
        if name:
            return self.search_entities(query=name, top_k=limit)

        return self.get_entities(entity_type=entity_type, limit=limit)

    def get_relationships_directional(
        self,
        entity_id: str,
        direction: str = "both",
        relation_type: Optional[str] = None,
    ) -> List[Relationship]:
        """
        Get relationships for an entity with direction control.

        This method provides compatibility with the existing GraphStore API.

        Args:
            entity_id: The entity ID.
            direction: "outgoing", "incoming", or "both".
            relation_type: Optional filter by relationship type.

        Returns:
            List of matching relationships.
        """
        if direction == "outgoing":
            pattern = "(e)-[r]->(other)"
        elif direction == "incoming":
            pattern = "(e)<-[r]-(other)"
        else:
            pattern = "(e)-[r]-(other)"

        type_filter = f":{relation_type}" if relation_type else ""

        query = f"""
        MATCH (e:Entity {{id: $entity_id}}){pattern.replace("[r]", f"[r{type_filter}]")}
        RETURN r.id as id, e.id as source_id, other.id as target_id,
               type(r) as relation_type, r.properties as properties,
               r.confidence as confidence, r.created_at as created_at
        """

        results = self._run_query(query, {"entity_id": entity_id})
        relationships = []
        for r in results:
            rel_id = (
                r["id"] or f"{r['source_id']}-{r['relation_type']}-{r['target_id']}"
            )
            relationships.append(
                Relationship(
                    id=rel_id,
                    source_id=r["source_id"],
                    target_id=r["target_id"],
                    relation_type=r["relation_type"],
                    properties=json.loads(r["properties"]) if r["properties"] else {},
                    confidence=r["confidence"] or 1.0,
                    created_at=(
                        datetime.fromisoformat(r["created_at"])
                        if r["created_at"]
                        else datetime.now(timezone.utc)
                    ),
                )
            )
        return relationships
