"""
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 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 intentionally avoid using timestamps as the signature, since that makes
        every poll look like a modification.
        """
        if not self._stata_client or not hasattr(self._stata_client, "stata"):
            return ""
        try:
            # Capture output so we can hash it deterministically.
            resp = self._stata_client.run_command_structured(f"graph describe {graph_name}", echo=False)
            if resp.success and resp.stdout:
                return resp.stdout
            if resp.error and resp.error.snippet:
                return resp.error.snippet
        except Exception:
            return ""
        return ""
    
    def _detect_graphs_via_pystata(self) -> List[str]:
        """Detect newly created graphs using direct pystata state access."""
        if not self._stata_client:
            return []
        
        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 only.
                    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()
            else:
                # Fallback to sfi Macro interface - only if stata is available
                if self._stata_client and hasattr(self._stata_client, 'stata'):
                    try:
                        from sfi import Macro
                        self._stata_client.stata.run("quietly graph dir, memory")
                        self._stata_client.stata.run("global mcp_graph_list `r(list)'")
                        graph_list_str = Macro.getGlobal("mcp_graph_list")
                        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()
            
            for graph_name in current_graphs:
                try:
                    signature = self._describe_graph_signature(graph_name)
                    state_info = {
                        "name": graph_name,
                        "exists": True,
                        "valid": bool(signature),
                        "signature": signature,
                    }

                    # Only update timestamps when the signature changes.
                    prev = self._last_graph_state.get(graph_name)
                    if prev is None or prev.get("signature") != signature:
                        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": ""}
            
        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:
            # Get current graph state via SFI
            current_graphs = set(self._get_current_graphs_from_pystata())
            
            # 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 for next comparison (stable signatures)
            new_state: Dict[str, Any] = {}
            for graph in current_graphs:
                sig = self._describe_graph_signature(graph)
                new_state[graph] = {
                    "name": graph,
                    "exists": True,
                    "valid": bool(sig),
                    "signature": sig,
                    "timestamp": time.time(),
                }
            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
        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 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
                    for callback in self._cache_callbacks:
                        try:
                            callback(graph_name, success)
                        except Exception as e:
                            logger.warning(f"Cache callback failed for {graph_name}: {e}")
                
                except Exception as e:
                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
                    # Still notify callbacks of failure
                    for callback in self._cache_callbacks:
                        try:
                            callback(graph_name, False)
                        except Exception:
                            pass
        
        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
                    for callback in self._cache_callbacks:
                        try:
                            callback(graph_name, success)
                        except Exception as e:
                            logger.warning(f"Cache callback failed for {graph_name}: {e}")
                
                except Exception as e:
                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
                    # Still notify callbacks of failure
                    for callback in self._cache_callbacks:
                        try:
                            callback(graph_name, False)
                        except Exception:
                            pass
        
        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()
