# Gemini Functions

import gradio as gr
import os
import google.generativeai as genai
import json

def initialize_gemini():
    try:
        genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
        generation_config = {
            "temperature": 1,
            "top_p": 0.95,
            "top_k": 40,
            "max_output_tokens": 8192,
            "response_mime_type": "text/plain",
        }
        model = genai.GenerativeModel(
            # model_name="gemini-1.5-pro",
            model_name="gemini-2.0-flash-exp",
            generation_config=generation_config,
        )
        return model
    except Exception as e:
        raise gr.Error(f"Error initializing Gemini: {str(e)}")

def create_prompt(extracted_text: str, path_to_data_to_extract: str) -> str:

    # load data to extract
    with open(path_to_data_to_extract, 'r', encoding='utf-8') as file:
        data_to_extract = json.load(file)

    prompt = f"""Tu es un assistant juridique expert en analyse de documents judiciaires français.
    Je vais te fournir le contenu d'un document judiciaire extrait d'un PDF.
    Ta tâche est d'analyser ce texte et d'en extraire les informations suivantes de manière précise :

    {json.dumps(data_to_extract, indent=2, ensure_ascii=False)}

    Voici quelques règles à suivre :
    - Si une information n'est pas présente dans le texte, indique "Non spécifié" pour cette catégorie.
    - Pour les noms des parties (demandeurs et défendeurs, et leurs avocats), liste tous ceux que tu trouves
    - Assure-toi de différencier correctement les demandeurs des défendeurs.
    - Si tu n'es pas sûr d'une information, indique-le clairement.

    Présente tes résultats sous forme de JSON, en utilisant les catégories mentionnées ci-dessus.

    Voici le contenu du document :

    {extracted_text.strip()}

    Analyse ce texte et fournis-moi les informations demandées au format JSON uniquement.""".strip()

    return prompt

def extract_data_with_gemini(text_file_path: str, path_to_data_to_extract: str) -> dict:
    try:
        # Initialize Gemini
        model = initialize_gemini()

        # Read the extracted text
        with open(text_file_path, 'r', encoding='utf-8') as f:
            extracted_text = f.read()

        # Create prompt and get response
        prompt = create_prompt(extracted_text, path_to_data_to_extract)
        response = model.generate_content(prompt)

        # Parse the JSON response
        try:
            # Extract JSON from the response text
            json_str = response.text
            if "json" in json_str.lower():
                json_str = json_str.split("json")[1].split("```")[0]
            elif "```" in json_str:
                json_str = json_str.split("```")[1]
            result = json.loads(json_str)
        except:
            result = {"error": "Failed to parse JSON response", "raw_response": response.text}

        return result
    except Exception as e:
        raise gr.Error(f"Error in Gemini processing: {str(e)}")