import google.generativeai as genai
import os
import json
from dotenv import load_dotenv
load_dotenv()

api_key = os.getenv("GOOGLE_API_KEY")
if not api_key:
    raise ValueError("GOOGLE_API_KEY environment variable is not set. Please add it to your .env file")

print(f"Google API Key loaded: {api_key[:10]}..." if api_key else "No API key found")
genai.configure(api_key=api_key)

def query_gemini(questions, contexts):
    try:
        context = "\n\n".join(contexts)
        
        # Create a numbered list of questions
        questions_text = "\n".join([f"{i+1}. {q}" for i, q in enumerate(questions)])
        
        prompt = f"""You are an insurance policy assistant. Based on the below document snippets, answer the following questions precisely.

IMPORTANT INSTRUCTIONS:
1. Only respond based on the context provided. If information is not found in the context, respond with "Not Found".
2. Provide clear, concise answers that directly address each question.
3. Return your response in the exact JSON format shown below.
4. Give complete, informative responses based on the provided context.
5. Answer each question in the order provided.

Context:
{context}

Questions:
{questions_text}

Return your response in this exact JSON format:
{{
    "answers": [
        "Answer to question 1",
        "Answer to question 2",
        "Answer to question 3",
        ...
    ]
}}

Ensure each answer is comprehensive and directly addresses the corresponding question. If information is not found in the context for any question, respond with "Not Found" for that question."""

        model = genai.GenerativeModel('gemini-2.0-flash-exp')
        response = model.generate_content(prompt)
        response_text = response.text.strip()
        
        # Try to parse the response as JSON
        try:
            # Remove any markdown code blocks if present
            if response_text.startswith("```json"):
                response_text = response_text.replace("```json", "").replace("```", "").strip()
            elif response_text.startswith("```"):
                response_text = response_text.replace("```", "").strip()
            
            parsed_response = json.loads(response_text)
            return parsed_response
        except json.JSONDecodeError:
            # If JSON parsing fails, return a structured response
            print(f"Failed to parse JSON response: {response_text}")
            return {"answers": ["Error parsing response"] * len(questions)}
            
    except Exception as e:
        print(f"Error in query_gemini: {str(e)}")
        return {"answers": [f"Error generating response: {str(e)}"] * len(questions)}