# utils/pandasai_setup.py
import os
import logging
import pandasai as pai
from pandasai_litellm import LiteLLM # Ensure this import matches your installed library

# Configure logger for this module
logger = logging.getLogger(__name__)

# It's good practice to define constants at the top or in a config file
DEFAULT_PANDASAI_MODEL = "gemini/gemini-2.5-flash-preview-05-20" # Using a common default

def configure_pandasai(api_key: str, model_name: str = None):
    """
    Configures PandasAI with LiteLLM using the provided API key and model.

    Args:
        api_key: The Google API key.
        model_name: The specific model to use (e.g., "gemini/gemini-1.5-flash-latest"). 
                    If None, uses DEFAULT_PANDASAI_MODEL.
    """
    if not api_key:
        logger.error("PandasAI Configuration Error: API key is missing.")
        # Depending on strictness, you might raise an error or just log
        # raise ValueError("API key must be provided for PandasAI configuration")
        return

    os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "False"
    os.environ["GOOGLE_API_KEY"] = api_key


    selected_model = model_name if model_name else DEFAULT_PANDASAI_MODEL
    
    try:
        llm = LiteLLM(
            model=DEFAULT_PANDASAI_MODEL , # Use the selected model
            api_key=api_key
        )

        # PandasAI configuration
        pai.config.set({
            "llm": llm,
            "temperature": 0.3,  # Lower temperature for more consistent results
            "max_retries": 3
        })
        logger.info(f"PandasAI configured successfully with model: {selected_model}")
        logger.info(f"PandasAI LLM object: {llm}")


    except ImportError:
        logger.error("PandasAI or pandasai_litellm is not installed. Please install the required packages.")
    except Exception as e:
        logger.error(f"Error configuring PandasAI: {e}", exc_info=True)