# 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

    selected_model = model_name if model_name else DEFAULT_PANDASAI_MODEL
    
    try:
        llm = LiteLLM(
            model=selected_model, # Use the selected model
            api_key=api_key
            # You might need to add other parameters for LiteLLM depending on the provider
            # e.g., if not using a Google model directly via gemini provider in LiteLLM
        )

        # PandasAI configuration
        pai.config.set({
            "llm": llm,
            "verbose": os.environ.get("PANDASAI_VERBOSE", "False").lower() == "true",
            "enable_cache": True,
            "enforce_privacy": False, # Be cautious with this in production
            "save_charts": False, # Set to True if you want to save charts locally
            # "save_charts_path": "charts_output", # Define path if saving charts
            "custom_whitelisted_dependencies": [], # Add any custom dependencies if needed
            "max_retries": 3, # Default retries for PandasAI operations
            "temperature": 0.3, # Lower temperature for more deterministic/factual outputs
            # "open_charts": False # Whether to automatically open charts
        })
        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)