"""
MCALStore: LangGraph BaseStore implementation with goal-aware search.

This is the main integration point for LangGraph users.
"""

from __future__ import annotations

import asyncio
import json
import threading
from datetime import datetime, timezone, timedelta
from typing import Any, Dict, Iterable, List, Literal, Tuple, TYPE_CHECKING

from mcal_langgraph._compat import (
    LANGGRAPH_AVAILABLE,
    check_langgraph,
    BaseStore,
    Item,
    SearchItem,
    GetOp,
    PutOp,
    SearchOp,
    ListNamespacesOp,
    NamespacePath,
)

if TYPE_CHECKING:
    from mcal import MCAL

# Type aliases
Op = Any
Result = Any


def _run_sync(coro):
    """
    Run async coroutine synchronously, handling existing event loops.
    
    Works in Lambda, Jupyter, FastAPI, and other async environments.
    """
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        # No running loop - safe to use asyncio.run()
        return asyncio.run(coro)
    else:
        # Loop already running - run in thread pool
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            future = pool.submit(asyncio.run, coro)
            return future.result()


class MCALStore(BaseStore if LANGGRAPH_AVAILABLE else object):
    """
    LangGraph-compatible store backed by MCAL.
    
    Implements LangGraph's BaseStore interface, providing MCAL's
    goal-aware memory capabilities through the standard LangGraph API.
    
    Key Advantages over generic stores:
    - Goal-aware search (MCAL's unique value)
    - Decision tracking in search results
    - Intent preservation across conversations
    
    Usage:
        from mcal import MCAL
        from mcal_langgraph import MCALStore
        
        mcal = MCAL(llm_provider="anthropic")
        store = MCALStore(mcal)
        
        # Use with LangGraph
        from langgraph.prebuilt import create_react_agent
        agent = create_react_agent(model, tools, store=store)
        
        # Or use directly
        await store.aput(("user_123",), "mem1", {"text": "Use PostgreSQL"})
        results = await store.asearch(("user_123",), query="database choice")
    """
    
    def __init__(self, mcal: "MCAL"):
        """
        Initialize MCALStore.
        
        Args:
            mcal: Initialized MCAL instance
        """
        check_langgraph()
        self._mcal = mcal
        self._lock = threading.RLock()  # Thread safety for writes
        
        # In-memory storage for items
        self._items: Dict[Tuple[str, ...], Dict[str, Dict[str, Any]]] = {}
        self._created_at: Dict[Tuple[str, ...], Dict[str, datetime]] = {}
        self._updated_at: Dict[Tuple[str, ...], Dict[str, datetime]] = {}
        
        # TTL support (Issue #057)
        self._ttl: Dict[Tuple[str, ...], Dict[str, float]] = {}  # TTL in minutes
        self._expires_at: Dict[Tuple[str, ...], Dict[str, datetime]] = {}  # Expiration times
    
    @property
    def mcal(self) -> "MCAL":
        """Get the underlying MCAL instance."""
        return self._mcal
    
    # ========== TTL Support (Issue #057) ==========
    
    def _is_expired(self, namespace: Tuple[str, ...], key: str) -> bool:
        """Check if an item has expired based on TTL."""
        if namespace not in self._expires_at:
            return False
        if key not in self._expires_at.get(namespace, {}):
            return False
        
        expires_at = self._expires_at[namespace][key]
        return datetime.now(timezone.utc) > expires_at
    
    def _refresh_ttl(self, namespace: Tuple[str, ...], key: str) -> None:
        """Refresh TTL by updating expiration time."""
        if namespace in self._ttl and key in self._ttl.get(namespace, {}):
            ttl_minutes = self._ttl[namespace][key]
            self._expires_at[namespace][key] = datetime.now(timezone.utc) + timedelta(minutes=ttl_minutes)
    
    def _cleanup_ttl(self, namespace: Tuple[str, ...], key: str) -> None:
        """Clean up TTL tracking for a deleted item."""
        if namespace in self._ttl and key in self._ttl.get(namespace, {}):
            del self._ttl[namespace][key]
        if namespace in self._expires_at and key in self._expires_at.get(namespace, {}):
            del self._expires_at[namespace][key]
    
    # ========== Abstract Methods (Required by BaseStore) ==========
    
    def batch(self, ops: Iterable[Op]) -> List[Result]:
        """Execute multiple operations synchronously in a single batch."""
        return _run_sync(self.abatch(ops))
    
    async def abatch(self, ops: Iterable[Op]) -> List[Result]:
        """Execute multiple operations asynchronously in a single batch."""
        results = []
        for op in ops:
            if isinstance(op, GetOp):
                result = await self.aget(op.namespace, op.key)
                results.append(result)
            elif isinstance(op, PutOp):
                await self.aput(op.namespace, op.key, op.value)
                results.append(None)
            elif isinstance(op, SearchOp):
                result = await self.asearch(
                    op.namespace_prefix,
                    query=op.query,
                    filter=op.filter,
                    limit=op.limit,
                    offset=op.offset,
                )
                results.append(result)
            elif isinstance(op, ListNamespacesOp):
                prefix = None
                suffix = None
                if op.match_conditions:
                    for cond in op.match_conditions:
                        if cond.match_type == "prefix":
                            prefix = cond.path
                        elif cond.match_type == "suffix":
                            suffix = cond.path
                result = await self.alist_namespaces(
                    prefix=prefix,
                    suffix=suffix,
                    max_depth=op.max_depth,
                    limit=op.limit,
                    offset=op.offset,
                )
                results.append(result)
            else:
                results.append(None)
        return results
    
    # ========== Get Operations ==========
    
    def get(
        self,
        namespace: Tuple[str, ...],
        key: str,
        *,
        refresh_ttl: bool | None = None,
    ) -> Item | None:
        """Retrieve a single item synchronously."""
        return _run_sync(self.aget(namespace, key, refresh_ttl=refresh_ttl))
    
    async def aget(
        self,
        namespace: Tuple[str, ...],
        key: str,
        *,
        refresh_ttl: bool | None = None,
    ) -> Item | None:
        """
        Retrieve a single item asynchronously.
        
        Args:
            namespace: The namespace tuple for the item
            key: The unique key within the namespace
            refresh_ttl: If True, refresh the TTL on access. Default: True for items with TTL.
        
        Returns:
            Item if found and not expired, None otherwise
        """
        if namespace not in self._items:
            return None
        if key not in self._items[namespace]:
            return None
        
        # Check TTL expiration (lazy expiration)
        if self._is_expired(namespace, key):
            # Item has expired - delete it
            with self._lock:
                if namespace in self._items and key in self._items[namespace]:
                    del self._items[namespace][key]
                    self._cleanup_ttl(namespace, key)
            return None
        
        # Refresh TTL on access if requested (default: True for items with TTL)
        if refresh_ttl is not False:
            with self._lock:
                self._refresh_ttl(namespace, key)
        
        value = self._items[namespace][key]
        created = self._created_at.get(namespace, {}).get(key, datetime.now(timezone.utc))
        updated = self._updated_at.get(namespace, {}).get(key, datetime.now(timezone.utc))
        
        return Item(
            namespace=namespace,
            key=key,
            value=value,
            created_at=created,
            updated_at=updated,
        )
    
    # ========== Put Operations ==========
    
    def put(
        self,
        namespace: Tuple[str, ...],
        key: str,
        value: Dict[str, Any],
        index: Literal[False] | List[str] | None = None,
        *,
        ttl: float | None = None,
    ) -> None:
        """Store or update an item synchronously."""
        _run_sync(self.aput(namespace, key, value, index, ttl=ttl))
    
    async def aput(
        self,
        namespace: Tuple[str, ...],
        key: str,
        value: Dict[str, Any],
        index: Literal[False] | List[str] | None = None,
        *,
        ttl: float | None = None,
    ) -> None:
        """
        Store or update an item asynchronously.
        
        Args:
            namespace: The namespace tuple for the item
            key: The unique key within the namespace
            value: The value to store (must be a dict)
            index: Fields to index for search. False to skip indexing.
            ttl: Time-to-live in minutes. None for no expiration.
        """
        now = datetime.now(timezone.utc)
        
        # Thread-safe write
        with self._lock:
            if namespace not in self._items:
                self._items[namespace] = {}
                self._created_at[namespace] = {}
                self._updated_at[namespace] = {}
            
            if key not in self._items[namespace]:
                self._created_at[namespace][key] = now
            
            self._items[namespace][key] = value
            self._updated_at[namespace][key] = now
            
            # Store TTL if provided (Issue #057)
            if ttl is not None:
                if namespace not in self._ttl:
                    self._ttl[namespace] = {}
                    self._expires_at[namespace] = {}
                self._ttl[namespace][key] = ttl
                self._expires_at[namespace][key] = now + timedelta(minutes=ttl)
        
        # Process through MCAL for goal-aware retrieval (outside lock)
        user_id = namespace[0] if namespace else "default"
        text = value.get("text", value.get("content", json.dumps(value)))
        
        if text and index is not False:
            try:
                await self._mcal.add(
                    messages=[{"role": "user", "content": text}],
                    user_id=user_id,
                    metadata={"key": key, "namespace": "/".join(namespace)},
                )
            except Exception:
                pass  # Don't fail if MCAL processing fails
    
    # ========== Delete Operations ==========
    
    def delete(self, namespace: Tuple[str, ...], key: str) -> None:
        """Delete an item synchronously."""
        _run_sync(self.adelete(namespace, key))
    
    async def adelete(self, namespace: Tuple[str, ...], key: str) -> None:
        """Delete an item asynchronously."""
        with self._lock:
            if namespace in self._items and key in self._items[namespace]:
                del self._items[namespace][key]
                if namespace in self._created_at and key in self._created_at[namespace]:
                    del self._created_at[namespace][key]
                if namespace in self._updated_at and key in self._updated_at[namespace]:
                    del self._updated_at[namespace][key]
                # Clean up TTL tracking
                self._cleanup_ttl(namespace, key)
    
    # ========== Search Operations ==========
    
    def search(
        self,
        namespace_prefix: Tuple[str, ...],
        /,
        *,
        query: str | None = None,
        filter: Dict[str, Any] | None = None,
        limit: int = 10,
        offset: int = 0,
        refresh_ttl: bool | None = None,
    ) -> List[SearchItem]:
        """Search for items synchronously."""
        return _run_sync(
            self.asearch(
                namespace_prefix,
                query=query,
                filter=filter,
                limit=limit,
                offset=offset,
                refresh_ttl=refresh_ttl,
            )
        )
    
    async def asearch(
        self,
        namespace_prefix: Tuple[str, ...],
        /,
        *,
        query: str | None = None,
        filter: Dict[str, Any] | None = None,
        limit: int = 10,
        offset: int = 0,
        refresh_ttl: bool | None = None,
    ) -> List[SearchItem]:
        """
        Search for items asynchronously.
        
        MCAL Enhancement: Uses goal-aware retrieval when query is provided,
        enriching results with goals and decisions.
        """
        results = []
        
        # If query provided, use MCAL's goal-aware search
        if query:
            user_id = namespace_prefix[0] if namespace_prefix else "default"
            try:
                mcal_results = await self._mcal.search(
                    query=query,
                    user_id=user_id,
                    limit=limit,
                )
                
                # Goals/decisions are on the SearchResult, not per-memory
                goals = [
                    str(g) for g in (mcal_results.active_goals if mcal_results else [])
                ]
                decisions = [
                    str(d) for d in (mcal_results.relevant_decisions if mcal_results else [])
                ]
                
                for i, mem in enumerate(mcal_results.memories if mcal_results else []):
                    if i < offset:
                        continue
                    if len(results) >= limit:
                        break
                    
                    now = datetime.now(timezone.utc)
                    results.append(SearchItem(
                        namespace=namespace_prefix,
                        key=f"mcal_{i}",
                        value={
                            "text": mem.content,
                            "goals": goals,
                            "decisions": decisions,
                            "score": mem.score or 0.0,
                            "metadata": mem.metadata or {},
                        },
                        created_at=now,
                        updated_at=now,
                        score=mem.score or 0.0,
                    ))
            except Exception:
                pass
        
        # Also search in-memory items
        for ns, items in list(self._items.items()):
            if not self._namespace_matches_prefix(ns, namespace_prefix):
                continue
            
            for key, value in list(items.items()):
                # Skip expired items (lazy expiration)
                if self._is_expired(ns, key):
                    with self._lock:
                        if ns in self._items and key in self._items[ns]:
                            del self._items[ns][key]
                            self._cleanup_ttl(ns, key)
                    continue
                
                if filter and not self._matches_filter(value, filter):
                    continue
                
                if query and not self._text_matches_query(value, query):
                    continue
                
                # Refresh TTL on search access if requested
                if refresh_ttl is not False:
                    with self._lock:
                        self._refresh_ttl(ns, key)
                
                created = self._created_at.get(ns, {}).get(key, datetime.now(timezone.utc))
                updated = self._updated_at.get(ns, {}).get(key, datetime.now(timezone.utc))
                
                results.append(SearchItem(
                    namespace=ns,
                    key=key,
                    value=value,
                    created_at=created,
                    updated_at=updated,
                    score=1.0 if query else None,
                ))
        
        return results[offset:offset + limit]
    
    # ========== List Namespaces ==========
    
    def list_namespaces(
        self,
        *,
        prefix: NamespacePath | None = None,
        suffix: NamespacePath | None = None,
        max_depth: int | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> List[Tuple[str, ...]]:
        """List namespaces synchronously."""
        return _run_sync(
            self.alist_namespaces(
                prefix=prefix,
                suffix=suffix,
                max_depth=max_depth,
                limit=limit,
                offset=offset,
            )
        )
    
    async def alist_namespaces(
        self,
        *,
        prefix: NamespacePath | None = None,
        suffix: NamespacePath | None = None,
        max_depth: int | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> List[Tuple[str, ...]]:
        """List namespaces asynchronously."""
        namespaces = set()
        
        for ns in self._items.keys():
            if prefix and not self._namespace_matches_prefix(ns, prefix):
                continue
            
            if suffix and not self._namespace_matches_suffix(ns, suffix):
                continue
            
            if max_depth and len(ns) > max_depth:
                ns = ns[:max_depth]
            
            namespaces.add(ns)
        
        sorted_ns = sorted(namespaces)
        return sorted_ns[offset:offset + limit]
    
    # ========== Helper Methods ==========
    
    def _namespace_matches_prefix(
        self,
        namespace: Tuple[str, ...],
        prefix: Tuple[str, ...],
    ) -> bool:
        """Check if namespace starts with prefix."""
        if not prefix:
            return True
        if len(namespace) < len(prefix):
            return False
        return namespace[:len(prefix)] == prefix
    
    def _namespace_matches_suffix(
        self,
        namespace: Tuple[str, ...],
        suffix: Tuple[str, ...],
    ) -> bool:
        """Check if namespace ends with suffix."""
        if not suffix:
            return True
        if len(namespace) < len(suffix):
            return False
        return namespace[-len(suffix):] == suffix
    
    def _matches_filter(
        self,
        value: Dict[str, Any],
        filter: Dict[str, Any],
    ) -> bool:
        """Check if value matches filter criteria."""
        for key, expected in filter.items():
            if key not in value:
                return False
            actual = value[key]
            
            if isinstance(expected, dict):
                for op, val in expected.items():
                    if op == "$eq" and actual != val:
                        return False
                    elif op == "$ne" and actual == val:
                        return False
                    elif op == "$gt" and not (actual > val):
                        return False
                    elif op == "$gte" and not (actual >= val):
                        return False
                    elif op == "$lt" and not (actual < val):
                        return False
                    elif op == "$lte" and not (actual <= val):
                        return False
            else:
                if actual != expected:
                    return False
        return True
    
    def _text_matches_query(
        self,
        value: Dict[str, Any],
        query: str,
    ) -> bool:
        """Simple text matching for query."""
        query_lower = query.lower()
        for field in ["text", "content", "memory", "description"]:
            if field in value:
                if query_lower in str(value[field]).lower():
                    return True
        return False


__all__ = ["MCALStore"]
