import google.generativeai as genai
from pathlib import Path
import gradio as gr

# Set up the model
generation_config = {
  "temperature": 1,
  "top_p": 1,
  "top_k": 32,
  "max_output_tokens": 4096,
}

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"
  }
]

# Configure the model with the API key
genai.configure(api_key = "AIzaSyA98JZGFwwaXr4LuuRMVDXhstasYU-XiME")

# Initialize the generative model
model = genai.GenerativeModel(model_name = "gemini-pro-vision",
                              generation_config = generation_config,
                              safety_settings = safety_settings)

def input_image_setup(file_loc):
    if not (img := Path(file_loc)).exists():
        raise FileNotFoundError(f"Could not find image: {img}")

    image_parts = [
        {
            "mime_type": "image/jpeg",
            "data": Path(file_loc).read_bytes()
        }
    ]
    return image_parts

def generate_gemini_response(input_prompt, text_input, image_loc):
    image_prompt = input_image_setup(image_loc)
    prompt_parts = [input_prompt + text_input, image_prompt[0]]
    response = model.generate_content(prompt_parts)
    return response.text

# Updated input prompt for car identification and specification
input_prompt = """This image shows a car.
Please analyze the image along with the prompted message and identify the make, model, and specialty of the car.
Based on your analysis, provide detailed specifications of the car.
The prompted message is """

def upload_file(files, text_input):
    file_paths = [file.name for file in files]
    if file_paths:
        response = generate_gemini_response(input_prompt, text_input, file_paths[0])
    return file_paths[0], response

# Setting up the Gradio interface
with gr.Blocks() as demo:
    header = gr.Label("Upload an image of a car and provide some details about it to get the car's specifications:")
    text_input = gr.Textbox(label="Provide any additional details or context about the car")
    image_output = gr.Image()
    upload_button = gr.UploadButton("Click to upload an image of the car", file_types=["image"], file_count="multiple")

    file_output = gr.Textbox(label="Car Specifications")
    combined_output = [image_output, file_output]

    upload_button.upload(upload_file, [upload_button, text_input], combined_output)

demo.launch(debug=True, share=True)