import os
import logging
import asyncio
import time
import json
import hashlib
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Union, Any, Tuple
from dataclasses import dataclass, field

try:
    from pymongo import MongoClient, errors
    from pymongo.collection import Collection
    from pymongo.database import Database
except ImportError:
    raise ImportError("Please install pymongo: pip install pymongo")

try:
    from motor.motor_asyncio import AsyncIOMotorClient
except ImportError:
    # Async support is optional but recommended
    AsyncIOMotorClient = None

@dataclass
class MongoGraphConfig:
    """Configuration class for MongoDB Graph helper."""
    host: str = "127.0.0.1"
    port: int = 27017
    username: str = ""
    password: str = ""
    database: str = "graph_db"
    node_collection: str = "graph_nodes"
    edge_collection: str = "graph_edges"
    uri: Optional[str] = None
    connect_timeout_ms: int = 5000
    socket_timeout_ms: int = 5000
    max_pool_size: int = 50
    enable_caching: bool = True
    cache_ttl: int = 3600
    log_level: str = "INFO"

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "MongoGraphConfig":
        """Create a config from a dictionary."""
        uri = data.get("uri") or data.get("host_url") or data.get("connection_string")
        host = data.get("host") or "127.0.0.1"
        port = int(data.get("port") or 27017)
        
        return cls(
            host=host,
            port=port,
            uri=uri,
            username=data.get("username") or data.get("user") or "",
            password=data.get("password") or data.get("pass") or "",
            database=data.get("database") or data.get("db_name") or "graph_db",
            node_collection=data.get("node_collection") or "graph_nodes",
            edge_collection=data.get("edge_collection") or "graph_edges",
            max_pool_size=data.get("max_pool_size", 50),
        )

@dataclass
class CacheEntry:
    """Cache entry with TTL support."""
    data: Any
    created_at: datetime
    ttl: int

    def is_expired(self) -> bool:
        return datetime.now() > self.created_at + timedelta(seconds=self.ttl)

class MongoGraphHelper:
    """
    MongoDB Graph Helper for LangXchange.
    Implements a property graph model using two MongoDB collections: nodes and edges.
    """

    def __init__(self, config: Optional[MongoGraphConfig] = None):
        self.config = config or MongoGraphConfig()
        self._client: Optional[MongoClient] = None
        self._async_client: Optional[Any] = None
        self._db: Optional[Database] = None
        self._async_db: Optional[Any] = None
        self._setup_logging()
        self._cache: Dict[str, CacheEntry] = {}

    def _setup_logging(self):
        logging.basicConfig(
            level=getattr(logging, self.config.log_level.upper(), logging.INFO),
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        )
        self.logger = logging.getLogger(self.__class__.__name__)

    def _get_connection_uri(self) -> str:
        if self.config.uri:
            return self.config.uri
        
        if self.config.username and self.config.password:
            return f"mongodb://{self.config.username}:{self.config.password}@{self.config.host}:{self.config.port}/{self.config.database}"
        return f"mongodb://{self.config.host}:{self.config.port}/{self.config.database}"

    def connect(self) -> "MongoGraphHelper":
        """Establish synchronous connection to MongoDB."""
        if self._client is None:
            uri = self._get_connection_uri()
            try:
                self._client = MongoClient(
                    uri,
                    connectTimeoutMS=self.config.connect_timeout_ms,
                    socketTimeoutMS=self.config.socket_timeout_ms,
                    maxPoolSize=self.config.max_pool_size
                )
                # Test connection
                self._client.admin.command('ping')
                self._db = self._client[self.config.database]
                self._setup_indexes()
                self.logger.info(f"Connected to MongoDB Graph at {self.config.host}")
            except Exception as e:
                self.logger.error(f"Failed to connect to MongoDB: {e}")
                raise ConnectionError(f"Failed to connect to MongoDB: {e}")
        return self

    async def connect_async(self) -> "MongoGraphHelper":
        """Establish asynchronous connection to MongoDB."""
        if self._async_client is None:
            if AsyncIOMotorClient is None:
                raise ImportError("motor is required for async support. Install with: pip install motor")
            
            uri = self._get_connection_uri()
            try:
                self._async_client = AsyncIOMotorClient(
                    uri,
                    connectTimeoutMS=self.config.connect_timeout_ms,
                    socketTimeoutMS=self.config.socket_timeout_ms,
                    maxPoolSize=self.config.max_pool_size
                )
                # Test connection
                await self._async_client.admin.command('ping')
                self._async_db = self._async_client[self.config.database]
                await self._setup_indexes_async()
                self.logger.info(f"Async connected to MongoDB Graph at {self.config.host}")
            except Exception as e:
                self.logger.error(f"Failed to async connect to MongoDB: {e}")
                raise ConnectionError(f"Failed to async connect to MongoDB: {e}")
        return self

    def close(self):
        """Close connections."""
        if self._client:
            self._client.close()
            self._client = None
            self._db = None
        self.logger.info("MongoDB Graph connections closed")

    async def close_async(self):
        """Close async connections."""
        if self._async_client:
            self._async_client.close()
            self._async_client = None
            self._async_db = None
        self.logger.info("Async MongoDB Graph connections closed")

    def _setup_indexes(self):
        """Setup necessary indexes for graph operations."""
        if self._db is not None:
            nodes = self._db[self.config.node_collection]
            edges = self._db[self.config.edge_collection]
            
            # Universal unique ID for nodes (could be 'id' or 'uri' or 'name')
            nodes.create_index("id", unique=False) # Not unique because multiple labels might share IDs in some models
            nodes.create_index([("id", 1), ("label", 1)], unique=True)
            
            # Edge indexes for traversal
            edges.create_index("from_id")
            edges.create_index("to_id")
            edges.create_index([("from_id", 1), ("to_id", 1), ("type", 1)], unique=True)

    async def _setup_indexes_async(self):
        """Setup necessary indexes for graph operations asynchronously."""
        if self._async_db is not None:
            nodes = self._async_db[self.config.node_collection]
            edges = self._async_db[self.config.edge_collection]
            
            await nodes.create_index([("id", 1), ("label", 1)], unique=True)
            await edges.create_index("from_id")
            await edges.create_index("to_id")
            await edges.create_index([("from_id", 1), ("to_id", 1), ("type", 1)], unique=True)

    # ---------------- Node Operations ----------------

    def merge_node(self, label: str, properties: Dict[str, Any], merge_on: Union[str, List[str]] = "id") -> Dict[str, Any]:
        """Merge a node into the database."""
        if self._db is None: self.connect()
        
        if isinstance(merge_on, str):
            merge_on = [merge_on]
            
        nodes = self._db[self.config.node_collection]
        
        # Build filter query for merging
        filter_query = {"label": label}
        for prop in merge_on:
            val = properties.get(prop)
            if val is not None:
                filter_query[prop] = val
        
        # If 'id' is in the unique index but missing in properties, default it to the first merge field or name
        if "id" not in properties:
             properties["id"] = properties.get(merge_on[0]) or properties.get("name")
        
        update_doc = {
            "$set": {**properties, "label": label, "updated_at": datetime.now()},
            "$setOnInsert": {"created_at": datetime.now()}
        }
        
        nodes.update_one(filter_query, update_doc, upsert=True)
        return properties

    async def merge_node_async(self, label: str, properties: Dict[str, Any], merge_on: Union[str, List[str]] = "id") -> Dict[str, Any]:
        """Merge a node into the database asynchronously."""
        if self._async_db is None: await self.connect_async()
        
        if isinstance(merge_on, str):
            merge_on = [merge_on]
            
        nodes = self._async_db[self.config.node_collection]
        
        # Build filter query for merging
        filter_query = {"label": label}
        for prop in merge_on:
            val = properties.get(prop)
            if val is not None:
                filter_query[prop] = val

        # If 'id' is in the unique index but missing in properties, default it to the first merge field or name
        if "id" not in properties:
             properties["id"] = properties.get(merge_on[0]) or properties.get("name")
        
        update_doc = {
            "$set": {**properties, "label": label, "updated_at": datetime.now()},
            "$setOnInsert": {"created_at": datetime.now()}
        }
        
        await nodes.update_one(filter_query, update_doc, upsert=True)
        return properties

    # ---------------- Relationship Operations ----------------

    def merge_relationship(
        self,
        source_label: str,
        source_match: Dict[str, Any],
        target_label: str,
        target_match: Dict[str, Any],
        rel_type: str,
        rel_properties: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Merge a relationship between two nodes."""
        if self._db is None: self.connect()
        
        edges = self._db[self.config.edge_collection]
        
        # Ensure ID exists in match
        src_id = source_match.get("id") or next(iter(source_match.values()))
        tgt_id = target_match.get("id") or next(iter(target_match.values()))
        
        # We don't necessarily enforce node existence here to match Cypher's MERGE behavior if nodes already exist,
        # but in a property graph on Mongo, we usually want to ensure we have the IDs.
        
        filter_query = {
            "from_id": src_id,
            "to_id": tgt_id,
            "type": rel_type
        }
        
        update_doc = {
            "$set": {
                **(rel_properties or {}),
                "from_label": source_label,
                "to_label": target_label,
                "updated_at": datetime.now()
            },
            "$setOnInsert": {"created_at": datetime.now()}
        }
        
        edges.update_one(filter_query, update_doc, upsert=True)
        return {"from_id": src_id, "to_id": tgt_id, "type": rel_type, "properties": rel_properties}

    async def merge_relationship_async(
        self,
        source_label: str,
        source_match: Dict[str, Any],
        target_label: str,
        target_match: Dict[str, Any],
        rel_type: str,
        rel_properties: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Merge a relationship between two nodes asynchronously."""
        if self._async_db is None: await self.connect_async()
        
        edges = self._async_db[self.config.edge_collection]
        
        src_id = source_match.get("id") or next(iter(source_match.values()))
        tgt_id = target_match.get("id") or next(iter(target_match.values()))
        
        filter_query = {
            "from_id": src_id,
            "to_id": tgt_id,
            "type": rel_type
        }
        
        update_doc = {
            "$set": {
                **(rel_properties or {}),
                "from_label": source_label,
                "to_label": target_label,
                "updated_at": datetime.now()
            },
            "$setOnInsert": {"created_at": datetime.now()}
        }
        
        await edges.update_one(filter_query, update_doc, upsert=True)
        return {"from_id": src_id, "to_id": tgt_id, "type": rel_type, "properties": rel_properties}

    # ---------------- Query and Traversal ----------------

    def query(self, pipeline: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Execute a custom aggregation pipeline (acts like a Cypher query)."""
        if self._db is None: self.connect()
        # Default query is against nodes collection
        return list(self._db[self.config.node_collection].aggregate(pipeline))

    async def query_async(self, pipeline: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Execute a custom aggregation pipeline asynchronously."""
        if self._async_db is None: await self.connect_async()
        cursor = self._async_db[self.config.node_collection].aggregate(pipeline)
        return await cursor.to_list(length=None)

    def get_neighbors(self, node_id: Any, direction: str = "out", max_depth: int = 1) -> List[Dict[str, Any]]:
        """
        Get neighbors of a node using $graphLookup.
        """
        if self._db is None: self.connect()
        
        # Start from the edges collection
        edges = self._db[self.config.edge_collection]
        
        if direction == "out":
            start_field = "from_id"
            connect_from = "to_id"
            connect_to = "from_id"
        elif direction == "in":
            start_field = "to_id"
            connect_from = "from_id"
            connect_to = "to_id"
        else: # both (not trivial in single $graphLookup without union)
            start_field = "from_id"
            connect_from = "to_id"
            connect_to = "from_id"

        pipeline = [
            {"$match": {start_field: node_id}},
            {
                "$graphLookup": {
                    "from": self.config.edge_collection,
                    "startWith": f"${connect_from}",
                    "connectFromField": connect_from,
                    "connectToField": connect_to,
                    "as": "traversed_edges",
                    "maxDepth": max_depth - 1
                }
            }
        ]
        
        return list(edges.aggregate(pipeline))

    def search_similar_by_vector(self, embedding: List[float], top_k: int = 10, index_name: str = "vector_index") -> List[Dict[str, Any]]:
        """
        Search for similar nodes using MongoDB Atlas Vector Search.
        Requires a Vector Search Index to be set up on the nodes collection.
        """
        if self._db is None: self.connect()
        
        pipeline = [
            {
                "$vectorSearch": {
                    "index": index_name,
                    "path": "embedding",
                    "queryVector": embedding,
                    "numCandidates": top_k * 10,
                    "limit": top_k
                }
            }
        ]
        
        return list(self._db[self.config.node_collection].aggregate(pipeline))

    async def search_similar_by_vector_async(self, embedding: List[float], top_k: int = 10, index_name: str = "vector_index") -> List[Dict[str, Any]]:
        """Search for similar nodes using MongoDB Atlas Vector Search asynchronously."""
        if self._async_db is None: await self.connect_async()
        
        pipeline = [
            {
                "$vectorSearch": {
                    "index": index_name,
                    "path": "embedding",
                    "queryVector": embedding,
                    "numCandidates": top_k * 10,
                    "limit": top_k
                }
            }
        ]
        
        cursor = self._async_db[self.config.node_collection].aggregate(pipeline)
        return await cursor.to_list(length=None)

    def __enter__(self):
        return self.connect()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
