from __future__ import annotations
"""
Graph creation detection for streaming Stata output.

This module provides functionality to detect when graphs are created
during Stata command execution and automatically cache them.
"""

import asyncio
import contextlib
import inspect
import re
import threading
import time
from typing import List, Set, Callable, Dict, Any
import logging


# SFI is always available
SFI_AVAILABLE = True

logger = logging.getLogger(__name__)


class GraphCreationDetector:
    """Detects graph creation using SFI-only detection with pystata integration."""
    
    def __init__(self, stata_client=None):
        self._lock = threading.Lock()
        self._detected_graphs: Set[str] = set()
        self._removed_graphs: Set[str] = set()
        self._unnamed_graph_counter = 0  # Track unnamed graphs for identification
        self._stata_client = stata_client
        self._last_graph_state: Dict[str, Any] = {}  # Track graph state changes

    def _describe_graph_signature(self, graph_name: str) -> str:
        """Return a stable signature for a graph.

        We avoid using Stata calls like 'graph describe' here because they are slow 
        (each call takes ~35ms) and would be called for every graph on every poll,
        bottlenecking the streaming output.
        
        Instead, we use name-based tracking tied to the Stata command execution 
        context. The signature is stable within a single command execution but 
        changes when a new command starts, allowing us to detect modifications 
        between commands without any Stata overhead.
        """
        if not self._stata_client:
            return ""
        
        # Access command_idx from stata_client if available
        cmd_idx = getattr(self._stata_client, "_command_idx", 0)
        # Always include command_idx to ensure modifications to named graphs are 
        # detected when a new command starts. Within a single command, the 
        # detected_graphs set in the cache will prevent duplicates.
        return f"{graph_name}_{cmd_idx}"

    def _get_graph_timestamp(self, graph_name: str) -> str:
        """Get the creation/modification timestamp of a graph using graph describe.
        
        The result is cached per command to optimize performance during streaming.
        """
        results = self._get_graph_timestamps([graph_name])
        return results.get(graph_name, "")

    def _get_graph_timestamps(self, graph_names: List[str]) -> Dict[str, str]:
        """Get timestamps for multiple graphs in a single Stata call to minimize overhead."""
        if not graph_names or not self._stata_client or not hasattr(self._stata_client, "stata"):
            return {}
        
        try:
            # Use the lock from client to prevent concurrency issues with pystata
            exec_lock = getattr(self._stata_client, "_exec_lock", None)
            ctx = exec_lock if exec_lock else contextlib.nullcontext()
            
            with ctx:
                hold_name = f"_mcp_detector_thold_{int(time.time() * 1000 % 1000000)}"
                self._stata_client.stata.run(f"capture _return hold {hold_name}", echo=False)
                try:
                    # Build a single Stata command to fetch all timestamps
                    stata_cmd = ""
                    for i, name in enumerate(graph_names):
                        resolved = self._stata_client._resolve_graph_name_for_stata(name)
                        stata_cmd += f"quietly graph describe {resolved}\n"
                        stata_cmd += f"macro define mcp_ts_{i} \"`r(command_date)'_`r(command_time)'\"\n"
                    
                    self._stata_client.stata.run(stata_cmd, echo=False)
                    
                    from sfi import Macro
                    results = {}
                    for i, name in enumerate(graph_names):
                        ts = Macro.getGlobal(f"mcp_ts_{i}")
                        if ts:
                            results[name] = ts
                    return results
                finally:
                    self._stata_client.stata.run(f"capture _return restore {hold_name}", echo=False)
        except Exception as e:
            logger.debug(f"Failed to get timestamps: {e}")
            return {}
    
    def _detect_graphs_via_pystata(self) -> List[str]:
        """Detect newly created graphs using direct pystata state access."""
        if not self._stata_client:
            return []
        
        with self._lock:
            try:
                # Get current graph state using pystata's sfi interface
                current_graphs = self._get_current_graphs_from_pystata()
                current_state = self._get_graph_state_from_pystata()
                
                # Compare with last known state to detect new graphs
                new_graphs = []
                
                # Check for new graph names
                for graph_name in current_graphs:
                    if graph_name not in self._last_graph_state and graph_name not in self._removed_graphs:
                        new_graphs.append(graph_name)
                
                # Check for state changes in existing graphs (modifications)
                for graph_name, state in current_state.items():
                    if graph_name in self._last_graph_state:
                        last_state = self._last_graph_state[graph_name]
                        # Compare stable signature.
                        if state.get("signature") != last_state.get("signature"):
                            if graph_name not in self._removed_graphs:
                                new_graphs.append(graph_name)
                
                # Update cached state
                self._last_graph_state = current_state.copy()
                
                return new_graphs
                
            except (ImportError, RuntimeError, ValueError, AttributeError) as e:
                # These are expected exceptions when SFI is not available or Stata state is inaccessible
                logger.debug(f"Failed to detect graphs via pystata (expected): {e}")
                return []
            except Exception as e:
                # Unexpected errors should be logged as errors
                logger.error(f"Unexpected error in pystata graph detection: {e}")
                return []
    
    def _get_current_graphs_from_pystata(self) -> List[str]:
        """Get current list of graphs using pystata's sfi interface."""
        try:
            # Use pystata to get graph list directly
            if self._stata_client and hasattr(self._stata_client, 'list_graphs'):
                return self._stata_client.list_graphs(force_refresh=True)
            else:
                # Fallback to sfi Macro interface - only if stata is available
                if self._stata_client and hasattr(self._stata_client, 'stata'):
                    # Access the lock from client to prevent concurrency issues with pystata
                    exec_lock = getattr(self._stata_client, "_exec_lock", None)
                    ctx = exec_lock if exec_lock else contextlib.nullcontext()
                    
                    with ctx:
                        try:
                            from sfi import Macro
                            hold_name = f"_mcp_detector_hold_{int(time.time() * 1000 % 1000000)}"
                            self._stata_client.stata.run(f"capture _return hold {hold_name}", echo=False)
                            try:
                                self._stata_client.stata.run("macro define mcp_graph_list \"\"", echo=False)
                                self._stata_client.stata.run("quietly graph dir, memory", echo=False)
                                self._stata_client.stata.run("macro define mcp_graph_list `r(list)'", echo=False)
                                graph_list_str = Macro.getGlobal("mcp_graph_list")
                            finally:
                                self._stata_client.stata.run(f"capture _return restore {hold_name}", echo=False)
                            return graph_list_str.split() if graph_list_str else []
                        except ImportError:
                            logger.warning("sfi.Macro not available for fallback graph detection")
                            return []
                else:
                    return []
        except Exception as e:
            logger.warning(f"Failed to get current graphs from pystata: {e}")
            return []
    
    def _get_graph_state_from_pystata(self) -> Dict[str, Any]:
        """Get detailed graph state information using pystata's sfi interface."""
        graph_state = {}
        
        try:
            current_graphs = self._get_current_graphs_from_pystata()
            cmd_idx = getattr(self._stata_client, "_command_idx", 0)
            
            # PRE-FETCH: Get timestamps for all current graphs in a single batch
            # to minimize Stata-Python boundary crossings.
            timestamps = self._get_graph_timestamps(current_graphs)
            
            for graph_name in current_graphs:
                try:
                    # Signature logic: 
                    # 1. Start with name+cmd_idx (fast)
                    fast_sig = self._describe_graph_signature(graph_name)
                    
                    # 2. If it's a new command for this graph, verify with timestamp
                    prev = self._last_graph_state.get(graph_name)
                    timestamp = timestamps.get(graph_name)
                    sig = fast_sig
                    
                    if prev and prev.get("cmd_idx") == cmd_idx:
                        # Already processed in this command. Keep the signature we decided on.
                        sig = prev.get("signature")
                    elif prev and prev.get("cmd_idx") != cmd_idx:
                        # Command jumped. We need to know if it's a REAL modification.
                        # We use the timestamp from graph describe.
                        if timestamp:
                            prev_ts = prev.get("timestamp_val")
                            if prev_ts and prev_ts == timestamp:
                                # Timestamp match! Reuse the OLD signature to avoid 
                                # reporting a modification just because cmd_idx jumped.
                                sig = prev.get("signature")
                            else:
                                # Timestamp changed or was missing. Use new fast_sig.
                                sig = fast_sig
                        else:
                            # Failed to get timestamp, fall back to fast_sig (safe default)
                            sig = fast_sig
                    
                    state_info = {
                        "name": graph_name,
                        "exists": True,
                        "valid": bool(sig),
                        "signature": sig,
                        "cmd_idx": cmd_idx,
                        "timestamp_val": timestamp,
                    }

                    # Only update visual timestamps when the signature changes.
                    if prev is None or prev.get("signature") != sig:
                        state_info["timestamp"] = time.time()
                    else:
                        state_info["timestamp"] = prev.get("timestamp", time.time())
                    
                    graph_state[graph_name] = state_info
                    
                except Exception as e:
                    logger.warning(f"Failed to get state for graph {graph_name}: {e}")
                    graph_state[graph_name] = {"name": graph_name, "timestamp": time.time(), "exists": False, "signature": "", "cmd_idx": cmd_idx}
            
        except Exception as e:
            logger.warning(f"Failed to get graph state from pystata: {e}")
        
        return graph_state
    
        
        
    def detect_graph_modifications(self, text: str = None) -> dict:
        """Detect graph modification/removal using SFI state comparison."""
        modifications = {"dropped": [], "renamed": [], "cleared": False}
        
        if not self._stata_client:
            return modifications
        
        try:
            # Use the more sophisticated state retrieval that handles timestamp verification
            new_state = self._get_graph_state_from_pystata()
            current_graphs = set(new_state.keys())
            
            # Compare with last known state to detect modifications
            if self._last_graph_state:
                last_graphs = set(self._last_graph_state.keys())
                
                # Detect dropped graphs (in last state but not current)
                dropped_graphs = last_graphs - current_graphs
                modifications["dropped"].extend(dropped_graphs)
                
                # Detect clear all (no graphs remain when there were some before)
                if last_graphs and not current_graphs:
                    modifications["cleared"] = True
            
            # Update last known state
            self._last_graph_state = new_state
            
        except Exception as e:
            logger.debug(f"SFI modification detection failed: {e}")
        
        return modifications
    
        
    def should_cache_graph(self, graph_name: str) -> bool:
        """Determine if a graph should be cached."""
        with self._lock:
            # Don't cache if already detected or removed
            if graph_name in self._detected_graphs or graph_name in self._removed_graphs:
                return False
            
            # Mark as detected
            self._detected_graphs.add(graph_name)
            return True
    
    def mark_graph_removed(self, graph_name: str) -> None:
        """Mark a graph as removed."""
        with self._lock:
            self._removed_graphs.add(graph_name)
            self._detected_graphs.discard(graph_name)
    
    def mark_all_cleared(self) -> None:
        """Mark all graphs as cleared."""
        with self._lock:
            self._detected_graphs.clear()
            self._removed_graphs.clear()
    
    def clear_detection_state(self) -> None:
        """Clear all detection state."""
        with self._lock:
            self._detected_graphs.clear()
            self._removed_graphs.clear()
            self._unnamed_graph_counter = 0
    
    def process_modifications(self, modifications: dict) -> None:
        """Process detected modifications."""
        with self._lock:
            # Handle dropped graphs
            for graph_name in modifications.get("dropped", []):
                self.mark_graph_removed(graph_name)
            
            # Handle renamed graphs
            for old_name, new_name in modifications.get("renamed", []):
                self.mark_graph_removed(old_name)
                self._detected_graphs.discard(new_name)  # Allow re-detection with new name
            
            # Handle clear all
            if modifications.get("cleared", False):
                self.mark_all_cleared()


class StreamingGraphCache:
    """Integrates graph detection with caching during streaming."""
    
    def __init__(self, stata_client, auto_cache: bool = False):
        self.stata_client = stata_client
        self.auto_cache = auto_cache
        # Use persistent detector from client if available, else create local one
        if hasattr(stata_client, "_graph_detector"):
            self.detector = stata_client._graph_detector
        else:
            self.detector = GraphCreationDetector(stata_client)
        self._lock = threading.Lock()
        self._cache_callbacks: List[Callable[[str, bool], None]] = []
        self._graphs_to_cache: List[str] = []
        self._cached_graphs: Set[str] = set()
        self._removed_graphs = set()  # Track removed graphs directly
        self._initial_graphs: Set[str] = set()  # Captured before execution starts
    
    def add_cache_callback(self, callback: Callable[[str, bool], None]) -> None:
        """Add callback for graph cache events."""
        with self._lock:
            self._cache_callbacks.append(callback)

    async def _notify_cache_callbacks(self, graph_name: str, success: bool) -> None:
        for callback in self._cache_callbacks:
            try:
                result = callback(graph_name, success)
                if inspect.isawaitable(result):
                    await result
            except Exception as e:
                logger.warning(f"Cache callback failed for {graph_name}: {e}")

    
    async def cache_detected_graphs_with_pystata(self) -> List[str]:
        """Enhanced caching method that uses pystata for real-time graph detection."""
        if not self.auto_cache:
            return []
        
        cached_names = []
        
        # First, try to get any newly detected graphs via pystata state
        if self.stata_client:
            try:
                # Get current state and check for new graphs

                pystata_detected = self.detector._detect_graphs_via_pystata()
                
                # Add any newly detected graphs to cache queue
                for graph_name in pystata_detected:
                    if graph_name not in self._cached_graphs and graph_name not in self._removed_graphs:
                        self._graphs_to_cache.append(graph_name)
                        
            except Exception as e:
                logger.warning(f"Failed to get pystata graph updates: {e}")
        
        # Process the cache queue
        with self._lock:
            graphs_to_process = self._graphs_to_cache.copy()
            self._graphs_to_cache.clear()
        
        # Get current graph list for verification
        try:
            current_graphs = self.stata_client.list_graphs()
        except Exception as e:
            logger.warning(f"Failed to get current graph list: {e}")
            return cached_names
        
        for graph_name in graphs_to_process:
            if graph_name in current_graphs and graph_name not in self._cached_graphs:
                try:
                    success = await asyncio.to_thread(self.stata_client.cache_graph_on_creation, graph_name)
                    if success:
                        cached_names.append(graph_name)
                        with self._lock:
                            self._cached_graphs.add(graph_name)
                    
                    # Notify callbacks
                    await self._notify_cache_callbacks(graph_name, success)
                
                except Exception as e:
                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
                    # Still notify callbacks of failure
                    await self._notify_cache_callbacks(graph_name, False)
        
        return cached_names
    
    async def cache_detected_graphs(self) -> List[str]:
        """Cache all detected graphs."""
        if not self.auto_cache:
            return []
        
        cached_names = []
        
        with self._lock:
            graphs_to_process = self._graphs_to_cache.copy()
            self._graphs_to_cache.clear()
        
        # Get current graph list for verification
        try:
            current_graphs = self.stata_client.list_graphs()
        except Exception as e:
            logger.warning(f"Failed to get current graph list: {e}")
            return cached_names
        
        for graph_name in graphs_to_process:
            if graph_name in current_graphs and graph_name not in self._cached_graphs:
                try:
                    success = await asyncio.to_thread(self.stata_client.cache_graph_on_creation, graph_name)
                    if success:
                        cached_names.append(graph_name)
                        with self._lock:
                            self._cached_graphs.add(graph_name)
                    
                    # Notify callbacks
                    await self._notify_cache_callbacks(graph_name, success)
                
                except Exception as e:
                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
                    # Still notify callbacks of failure
                    await self._notify_cache_callbacks(graph_name, False)
        
        return cached_names
    
    def get_cache_stats(self) -> dict:
        """Get caching statistics."""
        with self._lock:
            return {
                "auto_cache_enabled": self.auto_cache,
                "pending_cache_count": len(self._graphs_to_cache),
                "cached_graphs_count": len(self._cached_graphs),
                "detected_graphs_count": len(self.detector._detected_graphs),
                "removed_graphs_count": len(self.detector._removed_graphs),
            }
    
    def reset(self) -> None:
        """Reset the cache state."""
        with self._lock:
            self._graphs_to_cache.clear()
            self._cached_graphs.clear()
        self.detector.clear_detection_state()
