"""
Natural Language Query Engine for Outhad_Edge.

Main orchestrator that combines all AI components to provide natural language
query capabilities for eventstream analysis.

 Implementation:
- Integrates LLMAgent 
- Uses execute_with_llm_retry 
- Uses CodeValidator integration 
- Supports visualization and insights
- Enhanced caching with retrieve_context
- Production-ready with comprehensive error handling
"""

from typing import Any, Dict, Optional, Literal
import os
import warnings
import logging

logger = logging.getLogger(__name__)

# Load environment variables from .env file
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    # python-dotenv not installed, will rely on system environment variables
    pass

# Disable ChromaDB telemetry to avoid compatibility errors
os.environ.setdefault('ANONYMIZED_TELEMETRY', 'False')


class NaturalLanguageQuery:
    """
    Natural Language Query engine for Eventstream analysis.

    Combines LLM code generation, vector search, semantic caching, and safe execution
    to enable natural language queries over eventstream data.

    Example:
        >>> stream = Eventstream(df)
        >>> nlq = stream.nlq(model="gpt-4")
        >>> result = nlq.ask("What is the conversion rate?")
        >>> print(result['answer'])
    """

    def __init__(
        self,
        eventstream,
        model: Literal["gpt-4", "gpt-3.5-turbo", "claude-3-sonnet", "claude-3-opus"] = "gpt-4",
        temperature: float = 0.1,
        max_retries: int = 3,
        use_cache: bool = True,
        verbose: bool = False,
        enable_monitoring: bool = True,
    ):
        """
        Initialize Natural Language Query engine.

        Args:
            eventstream: Outhad_Edge Eventstream object
            model: LLM model to use
            temperature: LLM temperature (0-1, lower = more deterministic)
            max_retries: Maximum retries for code generation/execution
            use_cache: Enable semantic caching
            verbose: Print debug information
            enable_monitoring: Enable query monitoring and logging 
        """
        self.eventstream = eventstream
        self.model = model
        self.temperature = temperature
        self.max_retries = max_retries
        self.use_cache = use_cache
        self.verbose = verbose
        self.enable_monitoring = enable_monitoring

        # Initialize components
        self._initialize_components()

    def _initialize_components(self) -> None:
        """Initialize all AI components ."""
        from .semantic_layer import SemanticLayer
        from .vector_store import VectorStore
        from .cache_manager import CacheManager

        # Semantic layer for schema understanding 
        self.semantic_layer = SemanticLayer(self.eventstream)
        if self.verbose:
            logger.info("Initialized SemanticLayer")

        # Vector store for RAG 
        try:
            self.vector_store = VectorStore()
            if self.vector_store.count() == 0:
                from .vector_store import initialize_vector_store
                self.vector_store = initialize_vector_store()
            if self.verbose:
                logger.info(f"Initialized VectorStore ({self.vector_store.count()} examples)")
        except Exception as e:
            warnings.warn(f"Vector store initialization failed: {e}")
            self.vector_store = None

        # Code executor with  enhancements
        from .code_executor import CodeExecutor
        self.code_executor = CodeExecutor(
            timeout_seconds=60,
            use_validator=True,  # : Use CodeValidator
            strict_mode=True
        )
        if self.verbose:
            logger.info("Initialized CodeExecutor (with CodeValidator)")

        # Cache manager
        if self.use_cache:
            try:
                self.cache_manager = CacheManager(use_redis=True)
            except Exception:
                warnings.warn("Redis cache not available, using in-memory cache")
                self.cache_manager = CacheManager(use_redis=False)
        else:
            self.cache_manager = None

        # Initialize LLM Agent 
        self._initialize_llm_agent()

        # Initialize monitoring 
        if self.enable_monitoring:
            try:
                from .monitoring import NLQMonitoring
                self.monitor = NLQMonitoring(enable_console=self.verbose)
                if self.verbose:
                    logger.info("Initialized NLQMonitoring ")
            except Exception as e:
                warnings.warn(f"Monitoring initialization failed: {e}")
                self.monitor = None
        else:
            self.monitor = None

    def _initialize_llm_agent(self) -> None:
        """Initialize LLM Agent ."""
        try:
            from .llm_agent import LLMAgent

            # Use LLMAgent instead of direct LLM 
            self.llm_agent = LLMAgent(
                model=self.model,
                temperature=self.temperature,
                verbose=self.verbose
            )

            if self.verbose:
                logger.info(f"Initialized LLMAgent (model={self.model}, temp={self.temperature})")

        except Exception as e:
            raise RuntimeError(
                f"Failed to initialize LLM Agent: {e}\n"
                "Make sure API keys are set in .env file:\n"
                "  - OPENAI_API_KEY for GPT models\n"
                "  - ANTHROPIC_API_KEY for Claude models"
            )

    def ask(self, question: str, return_code: bool = False, explain: bool = True) -> Dict[str, Any]:
        """
        Ask a natural language question about the eventstream.

        Args:
            question: Natural language question
            return_code: If True, include generated code in response
            explain: If True, include insights and explanation in response 

        Returns:
            dict: Query result with 'answer', 'result', 'code' (optional), 'cached', 'insights' (optional)
        """
        import time
        start_time = time.time()

        if self.verbose:
            print(f"Question: {question}")

        # Check cache first
        cached = False
        if self.cache_manager:
            cached_result = self.cache_manager.get(question)
            if cached_result:
                if self.verbose:
                    print("✓ Retrieved from cache")
                cached_result["cached"] = True
                cached = True

                # Log cached query 
                if self.monitor:
                    latency = time.time() - start_time
                    self.monitor.log_query(
                        query=question,
                        success=True,
                        latency=latency,
                        model=self.model,
                        tokens_used=0,
                        cached=True,
                    )

                return cached_result

        # Get semantic context
        context = self.semantic_layer.get_full_context()

        # Retrieve similar examples using enhanced RAG 
        examples = ""
        if self.vector_store:
            try:
                # Use retrieve_context for better RAG ( enhancement)
                examples = self.vector_store.retrieve_context(question, n_results=2)
                if self.verbose:
                    print(f"✓ Retrieved {len(examples.split('Example'))-1 if examples else 0} similar examples")
            except Exception as e:
                if self.verbose:
                    print(f"Vector search failed: {e}")
                # Fallback to basic search
                try:
                    similar = self.vector_store.search(question, n_results=2)
                    examples = "\n\n".join(
                        [
                            f"Q: {ex['question']}\nCode:\n```python\n{ex['code']}\n```"
                            for ex in similar
                        ]
                    )
                except Exception:
                    examples = ""

        # Generate code
        code = self._generate_code(question, context, examples)

        if self.verbose:
            print(f"Generated code:\n{code}\n")

        # Execute code with retry
        exec_result = self._execute_with_retry(code, question)

        if not exec_result["success"]:
            # Log failed query 
            if self.monitor:
                latency = time.time() - start_time
                self.monitor.log_query(
                    query=question,
                    success=False,
                    latency=latency,
                    model=self.model,
                    tokens_used=0,
                    cached=False,
                    error=exec_result["error"],
                )

            return {
                "success": False,
                "error": exec_result["error"],
                "answer": f"Failed to answer: {exec_result['error']}",
                "question": question,
                "code": code if return_code else None,
                "cached": False,
            }

        # Format answer
        answer = self._format_answer(question, exec_result["result"], code)

        result = {
            "success": True,
            "answer": answer,
            "result": exec_result["result"],
            "question": question,
            "cached": False,
        }

        if return_code:
            result["code"] = code

        # Add insights and visualization support 
        if explain:
            insights = self._generate_insights(question, exec_result["result"], code)
            result["insights"] = insights
            result["visualization"] = self._suggest_visualization(exec_result["result"])

        # Cache result
        if self.cache_manager:
            self.cache_manager.set(question, result)

        # Log successful query 
        if self.monitor:
            latency = time.time() - start_time
            # Estimate tokens (rough approximation: ~4 chars per token)
            tokens_estimate = (len(question) + len(str(code)) + len(str(answer))) // 4
            self.monitor.log_query(
                query=question,
                success=True,
                latency=latency,
                model=self.model,
                tokens_used=tokens_estimate,
                cached=False,
            )

        return result

    def _generate_code(
        self, question: str, context: str, examples: str
    ) -> str:
        """Generate Python code to answer the question ( uses LLMAgent)."""
        try:
            # Use LLMAgent instead of direct LLM 
            code = self.llm_agent.generate_code(
                query=question,
                schema=context,
                examples=examples,
                glossary=""  # Can be enhanced later with domain glossary
            )
            return code

        except Exception as e:
            raise RuntimeError(f"Code generation failed: {e}")

    def _execute_with_retry(
        self, code: str, question: str
    ) -> Dict[str, Any]:
        """Execute code with retry and LLM-based error correction ( uses  method)."""
        # Use execute_with_llm_retry from  (integrates CodeValidator and LLMAgent)
        success, result, error = self.code_executor.execute_with_llm_retry(
            query=question,
            initial_code=code,
            stream=self.eventstream,
            llm_agent=self.llm_agent,
            max_retries=self.max_retries,
            verbose=self.verbose
        )

        if success:
            return result
        else:
            return {"success": False, "error": error}

    def _fix_code(self, code: str, error: str, question: str) -> str:
        """
        Use LLM to fix code error.

        Note: This method is deprecated in . Error correction is now handled
        by execute_with_llm_retry() which uses LLMAgent.correct_error().
        Kept for backward compatibility.
        """
        try:
            # Use LLMAgent for error correction 
            fixed_code = self.llm_agent.correct_error(
                query=question,
                code=code,
                error=error
            )
            return fixed_code

        except Exception as e:
            if self.verbose:
                print(f"Code correction failed: {e}")
            return code  # Return original code if correction fails

    def _format_answer(
        self, question: str, raw_result: Any, code: str
    ) -> str:
        """Format raw result into natural language answer."""
        # Simple formatting for now
        if raw_result is None:
            return "No result generated"

        if isinstance(raw_result, (int, float)):
            return f"{raw_result}"

        if isinstance(raw_result, str):
            return raw_result

        if hasattr(raw_result, "to_string"):
            return raw_result.to_string()

        return str(raw_result)

    def _generate_insights(self, question: str, result: Any, code: str) -> str:
        """
        Generate insights about the result .

        Args:
            question: User's query
            result: Execution result
            code: Generated code

        Returns:
            str: Human-readable insights
        """
        insights = []

        # Basic result type insight
        if isinstance(result, (int, float)):
            insights.append(f"The result is a numeric value: {result}")
        elif hasattr(result, 'shape'):
            # DataFrame or numpy array
            insights.append(f"Result shape: {result.shape}")
            if hasattr(result, 'columns'):
                insights.append(f"Columns: {', '.join(str(c) for c in result.columns)}")
        elif isinstance(result, (list, tuple)):
            insights.append(f"Result is a sequence with {len(result)} items")

        # Code complexity insight
        lines = code.strip().split('\n')
        non_empty_lines = [l for l in lines if l.strip() and not l.strip().startswith('#')]
        insights.append(f"Generated code: {len(non_empty_lines)} lines")

        return "\n".join(insights) if insights else "No additional insights available"

    def _suggest_visualization(self, result: Any) -> Optional[str]:
        """
        Suggest appropriate visualization type based on result .

        Args:
            result: Execution result

        Returns:
            str: Suggested visualization type or None
        """
        if isinstance(result, (int, float)):
            return "metric"  # Single metric display

        if hasattr(result, 'shape'):
            # DataFrame-like result
            if hasattr(result, 'columns'):
                num_cols = len(result.columns)
                num_rows = len(result)

                if num_rows <= 20:
                    return "table"  # Small enough to show as table
                elif num_cols == 2:
                    return "bar_chart"  # Two columns: category and value
                elif num_cols >= 3:
                    return "line_chart"  # Time series or multi-column data

        if isinstance(result, (list, tuple)):
            if len(result) <= 10:
                return "list"  # Simple list display
            else:
                return "table"  # Convert to table

        return None  # No specific visualization suggested

    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the NLQ engine."""
        stats = {
            "model": self.model,
            "temperature": self.temperature,
            "cache_enabled": self.use_cache,
            "vector_store_size": (
                self.vector_store.count() if self.vector_store else 0
            ),
        }

        if self.cache_manager:
            stats["cache_stats"] = self.cache_manager.get_stats()

        return stats

    def clear_cache(self) -> None:
        """Clear the semantic cache."""
        if self.cache_manager:
            self.cache_manager.clear()
            if self.verbose:
                print("Cache cleared")
