# chatbot_handler.py
import logging
import json
from google import genai
import os # For potential API key loading if Canvas injection fails for the library

# --- Gemini Configuration ---
# Option 1: Rely on Canvas to make this work with an empty key for the library.
# This is the preferred approach as per Canvas guidelines for 'fetch'.
GEMINI_API_KEY = "" 
GEMINI_API_KEY = os.getenv('GEMINI_API_KEY')

# Option 2: Fallback to environment variable if direct empty key doesn't work with the library via Canvas.
# if not GEMINI_API_KEY: # This check would be if we explicitly want to load from env
#     GEMINI_API_KEY = os.getenv('GEMINI_API_KEY_ENV_VAR_NAME') # Replace with your actual env var name if you use one
#     if not GEMINI_API_KEY:
#         logging.warning("GEMINI_API_KEY not found via direct assignment or environment variable.")
        # If you have a default key for local testing (NOT FOR PRODUCTION/CANVAS)
        # GEMINI_API_KEY = "YOUR_LOCAL_DEV_API_KEY" 

model = None
gen_config = None
safety_settings = []

try:
    if GEMINI_API_KEY is not None: # Check if it's set (even if empty string for Canvas)
        genai.configure(api_key=GEMINI_API_KEY)
        
        # As per general instructions, use gemini-2.0-flash if not told otherwise.
        MODEL_NAME = "gemini-2.0-flash" 
        model = genai.GenerativeModel(MODEL_NAME)
        
        gen_config = genai.types.GenerationConfig(
            temperature=0.7,
            top_k=1, # Per user's original config
            top_p=1, # Per user's original config
            max_output_tokens=2048, # Per user's original config
        )
        
        # Standard safety settings
        safety_settings = [
            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
        ]
        logging.info(f"Gemini model '{MODEL_NAME}' configured successfully.")
    else:
        logging.error("Gemini API Key is None. Model not configured.")

except Exception as e:
    logging.error(f"Failed to configure Gemini or instantiate model: {e}", exc_info=True)
    model = None # Ensure model is None if setup fails

def format_history_for_gemini(gradio_chat_history: list) -> list:
    """
    Converts Gradio chat history (list of dicts with 'role' and 'content')
    to Gemini API's 'contents' format (list of dicts with 'role' and 'parts').
    Gemini SDK expects roles 'user' and 'model'.
    """
    gemini_contents = []
    for msg in gradio_chat_history:
        # Map Gradio 'assistant' role to Gemini 'model' role
        role = "user" if msg["role"] == "user" else "model"
        
        content = msg.get("content")
        if isinstance(content, str):
            gemini_contents.append({"role": role, "parts": [{"text": content}]})
        else:
            logging.warning(f"Skipping non-string content in chat history for Gemini: {content}")
    return gemini_contents

async def generate_llm_response(user_message: str, plot_id: str, plot_label: str, chat_history_for_plot: list, plot_data_summary: str = None):
    """
    Generates a response from the LLM using Gemini API via the Python SDK.
    Args:
        user_message (str): The latest message from the user.
        plot_id (str): The ID of the plot being discussed.
        plot_label (str): The label of the plot being discussed.
        chat_history_for_plot (list): The current conversation history for this plot.
                                      This list ALREADY includes the initial assistant message (with data summary)
                                      and the latest user_message.
        plot_data_summary (str, optional): The textual summary of the plot data. 
                                           While it's in the history, passing it explicitly might be useful
                                           for future system prompt enhancements if needed.
    Returns:
        str: The LLM's response text.
    """
    logging.info(f"Generating LLM response for plot_id: {plot_id} ('{plot_label}'). User message: '{user_message}'")
    # Log the provided data summary for debugging
    # logging.debug(f"Data summary for '{plot_label}':\n{plot_data_summary}")


    if not model:
        logging.error("Gemini model not configured. Cannot generate LLM response.")
        return "I'm sorry, the AI model is not available at the moment. (Configuration Error)"

    # The chat_history_for_plot already contains the initial assistant message with the summary,
    # and the latest user message which triggered this call.
    gemini_formatted_history = format_history_for_gemini(chat_history_for_plot)

    if not gemini_formatted_history:
        logging.error("Cannot generate LLM response: Formatted history is empty or invalid.")
        return "I'm sorry, there was an issue processing the conversation history."

    # Optional: Construct a system instruction if desired, though the initial message in history helps.
    # system_instruction_text = (
    #     f"You are an expert in Employer Branding and LinkedIn social media strategy. "
    #     f"You are discussing the graph: '{plot_label}' (ID: '{plot_id}'). "
    #     f"A data summary for this graph was provided in your initial message: \n---\n{plot_data_summary}\n---\n"
    #     f"Refer to this summary and the conversation history to answer questions. "
    #     f"If specific data is not in the summary, clearly state that the provided snapshot doesn't contain that detail."
    # )
    # contents_for_api = [{"role": "system", "parts": [{"text": system_instruction_text}]}] + gemini_formatted_history
    # For now, relying on the summary being in the `gemini_formatted_history` via the first assistant message.

    try:
        logging.debug(f"Sending to Gemini API. History: {json.dumps(gemini_formatted_history, indent=2)}")
        
        response = await model.generate_content_async(
            contents=gemini_formatted_history, # History already includes user's latest message
            generation_config=gen_config,
            safety_settings=safety_settings
        )
        
        # logging.debug(f"LLM API Raw Response object for '{plot_label}': {response}")

        # Check for blocking based on prompt_feedback first (as per SDK examples)
        if response.prompt_feedback and response.prompt_feedback.block_reason:
            reason = response.prompt_feedback.block_reason.name # e.g., 'SAFETY'
            # safety_ratings_info = [f"{rating.category.name}: {rating.probability.name}" for rating in response.prompt_feedback.safety_ratings]
            # details = f" Safety Ratings: {', '.join(safety_ratings_info)}" if safety_ratings_info else ""
            logging.warning(f"Content blocked by API (prompt_feedback) for '{plot_label}'. Reason: {reason}.")
            return f"I'm sorry, your request was blocked due to content policy: {reason}."

        # Accessing response text (handle multi-part if any)
        if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
            response_text = "".join(part.text for part in response.candidates[0].content.parts)
            logging.info(f"LLM generated response for '{plot_label}': {response_text[:150]}...")
            return response_text
        else:
            # This case might occur if the response was empty but not blocked by prompt_feedback
            # (e.g. finish_reason other than SAFETY, or no candidates)
            finish_reason_str = "UNKNOWN"
            if response.candidates and response.candidates[0].finish_reason:
                 finish_reason_str = response.candidates[0].finish_reason.name # e.g. 'STOP', 'MAX_TOKENS', 'SAFETY', 'RECITATION', 'OTHER'
            
            if finish_reason_str == 'SAFETY': # Content blocked at candidate level
                 logging.warning(f"Content blocked by API (candidate safety) for '{plot_label}'. Finish Reason: {finish_reason_str}.")
                 return f"I'm sorry, I can't provide a response due to safety filters regarding: {finish_reason_str}."

            logging.error(f"Unexpected LLM API response structure or empty content for '{plot_label}'. Finish Reason: {finish_reason_str}. Full response: {response}")
            return f"Sorry, I received an unexpected or empty response from the AI model (Finish Reason: {finish_reason_str})."

    except google.api_core.exceptions.PermissionDenied as e:
        logging.error(f"LLM API Permission Denied (Status 403) for '{plot_label}': {e}", exc_info=True)
        return "Sorry, there's an issue with API permissions. Please ensure the API key is correct and the service is enabled. (Error 403)"
    except google.api_core.exceptions.InvalidArgument as e:
        logging.error(f"LLM API Invalid Argument (Status 400) for '{plot_label}': {e}. History: {json.dumps(gemini_formatted_history, indent=2)}", exc_info=True)
        return "Sorry, there was an issue with the request sent to the AI model (e.g. malformed history). (Error 400)"
    except google.api_core.exceptions.GoogleAPIError as e: # Catch other Google API errors
        logging.error(f"Google API Error during LLM call for '{plot_label}': {e}", exc_info=True)
        return f"An API error occurred while trying to get an AI response: {type(e).__name__}."
    except Exception as e:
        logging.error(f"Generic error during LLM call for '{plot_label}': {e}", exc_info=True)
        return f"An unexpected error occurred while trying to get an AI response: {type(e).__name__}."