"""
PandasAI adapter for qry-doc.

This module provides the PandasAIAdapter class that wraps PandasAI
functionality and adapts it to the qry-doc interface.
"""
from pathlib import Path
from typing import Any, Optional, Protocol, Union

import pandas as pd
from pandasai import SmartDataframe

from qry_doc.exceptions import QueryError
from qry_doc.validators import OutputValidator


class LLMProvider(Protocol):
    """
    Protocol for LLM providers.
    
    Any LLM that implements this interface can be used with qry-doc.
    Compatible with OpenAI, Anthropic, LiteLLM, and local models.
    """
    
    def __call__(self, prompt: str) -> str:
        """Generate a response for the given prompt."""
        ...


class PandasAIAdapter:
    """
    Adapter for PandasAI that provides a consistent interface for qry-doc.
    
    Features:
    - Query execution with type-specific returns
    - Prompt engineering for DataFrame/text responses
    - Chart path configuration
    - Error handling and sanitization
    """
    
    # Prompt suffixes for different output types
    DATAFRAME_SUFFIX = " Return the result strictly as a Pandas DataFrame object. Do not explain."
    TEXT_SUFFIX = " Provide a clear, conversational answer explaining the result."
    
    def __init__(
        self,
        df: pd.DataFrame,
        llm: Any,
        chart_save_path: Optional[Path] = None
    ) -> None:
        """
        Initialize the PandasAI adapter.
        
        Args:
            df: The DataFrame to query.
            llm: The LLM provider (OpenAI, Anthropic, etc.).
            chart_save_path: Optional path for saving generated charts.
        """
        self._df = df
        self._llm = llm
        self._chart_save_path = chart_save_path
        
        # Configure PandasAI
        config = {
            "llm": llm,
            "enable_cache": False,  # Disable cache for fresh results
            "open_charts": False,   # Don't open chart windows
        }
        
        if chart_save_path:
            config["save_charts"] = True
            config["save_charts_path"] = str(chart_save_path)
        
        self._sdf = SmartDataframe(df, config=config)
    
    def query(self, question: str) -> Any:
        """
        Execute a query and return the raw result.
        
        Args:
            question: Natural language question about the data.
            
        Returns:
            The raw result from PandasAI (could be DataFrame, string, number, etc.).
            
        Raises:
            QueryError: If the query fails.
        """
        try:
            result = self._sdf.chat(question)
            return result
        except Exception as e:
            sanitized = OutputValidator.sanitize_error_message(e)
            raise QueryError(
                user_message=f"Error al procesar la consulta: {sanitized}",
                internal_error=e
            )
    
    def query_as_dataframe(self, question: str) -> pd.DataFrame:
        """
        Execute a query and ensure the result is a DataFrame.
        
        Uses prompt engineering to encourage DataFrame output.
        
        Args:
            question: Natural language question about the data.
            
        Returns:
            A pandas DataFrame with the query results.
            
        Raises:
            QueryError: If the query fails or doesn't return a DataFrame.
        """
        # Enhance prompt to request DataFrame output
        enhanced_question = question + self.DATAFRAME_SUFFIX
        
        try:
            result = self._sdf.chat(enhanced_question)
            
            # Validate result is a DataFrame
            if isinstance(result, pd.DataFrame):
                return result
            
            # Try to convert to DataFrame if possible
            if isinstance(result, (list, dict)):
                try:
                    return pd.DataFrame(result)
                except Exception:
                    pass
            
            # If we got a single value, wrap it
            if isinstance(result, (int, float, str)):
                return pd.DataFrame({"result": [result]})
            
            raise QueryError(
                user_message=(
                    "La consulta no retornó datos tabulares. "
                    "Intente reformular la pregunta para obtener una tabla de datos."
                )
            )
            
        except QueryError:
            raise
        except Exception as e:
            sanitized = OutputValidator.sanitize_error_message(e)
            raise QueryError(
                user_message=f"Error al extraer datos: {sanitized}",
                internal_error=e
            )
    
    def query_as_text(self, question: str) -> str:
        """
        Execute a query and ensure the result is a text response.
        
        Uses prompt engineering to encourage conversational output.
        
        Args:
            question: Natural language question about the data.
            
        Returns:
            A string with the conversational response.
            
        Raises:
            QueryError: If the query fails.
        """
        # Enhance prompt to request conversational output
        enhanced_question = question + self.TEXT_SUFFIX
        
        try:
            result = self._sdf.chat(enhanced_question)
            
            # Convert result to string
            if isinstance(result, str):
                return result
            
            if isinstance(result, pd.DataFrame):
                # Summarize DataFrame as text
                if result.empty:
                    return "No se encontraron datos que coincidan con la consulta."
                
                # Create a simple text summary
                n_rows, n_cols = result.shape
                summary = f"Se encontraron {n_rows} registros con {n_cols} columnas.\n\n"
                
                # Add first few rows as text
                if n_rows <= 5:
                    summary += result.to_string(index=False)
                else:
                    summary += f"Primeros 5 registros:\n{result.head().to_string(index=False)}"
                
                return summary
            
            # Convert other types to string
            return str(result)
            
        except QueryError:
            raise
        except Exception as e:
            sanitized = OutputValidator.sanitize_error_message(e)
            raise QueryError(
                user_message=f"Error al procesar la consulta: {sanitized}",
                internal_error=e
            )
    
    def generate_chart(self, question: str) -> Optional[Path]:
        """
        Execute a query that generates a chart.
        
        Args:
            question: Natural language question requesting a visualization.
            
        Returns:
            Path to the generated chart file, or None if no chart was created.
            
        Raises:
            QueryError: If the query fails.
        """
        if not self._chart_save_path:
            raise QueryError(
                user_message="No se configuró una ruta para guardar gráficos."
            )
        
        # Enhance prompt to request chart
        chart_question = f"Create a chart or plot showing: {question}"
        
        try:
            self._sdf.chat(chart_question)
            
            # Look for generated chart
            chart_path = Path(self._chart_save_path)
            if chart_path.is_dir():
                # Find most recent chart file
                charts = list(chart_path.glob("*.png")) + list(chart_path.glob("*.jpg"))
                if charts:
                    charts.sort(key=lambda p: p.stat().st_mtime, reverse=True)
                    return charts[0]
            
            return None
            
        except QueryError:
            raise
        except Exception as e:
            sanitized = OutputValidator.sanitize_error_message(e)
            raise QueryError(
                user_message=f"Error al generar el gráfico: {sanitized}",
                internal_error=e
            )
    
    @property
    def dataframe(self) -> pd.DataFrame:
        """Get the underlying DataFrame."""
        return self._df
    
    @property
    def columns(self) -> list[str]:
        """Get the column names of the DataFrame."""
        return list(self._df.columns)
    
    @property
    def shape(self) -> tuple[int, int]:
        """Get the shape of the DataFrame (rows, columns)."""
        return self._df.shape
