import os
import uuid
from crewai import Agent, Task, Crew
from langchain_groq import ChatGroq
from PIL import Image
import gradio as gr

# Initialize the LLM
llm = ChatGroq(
    groq_api_key="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",
    model_name="llama3-70b-8192",
)

# Define the agent with a goal to analyze and describe images
image_caption_agent = Agent(
    role='Image Captioning Agent',
    goal='Analyze the provided image and generate a title and description.',
    backstory=( 
        "You are an Image Captioning Agent. Your job is to analyze and get what is inside the provided image and generate a relevant title and description."
    ),
    verbose=True,
    llm=llm,
)

# Ensure the folder for saving images exists
SAVE_DIR = './uploaded_images'
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR)

def process_image_with_agent(image_path, agent):
    try:
        image = Image.open(image_path)
    except Exception as e:
        return f"Error opening image: {e}"

    task_description = f"Analyze the image and generate a title and description. Image path: {image_path}"
    
    image_caption_task = Task(
        description=task_description,
        agent=agent,
        human_input=False,
        expected_output="Title and description generated by the agent"
    )

    crew = Crew(
        agents=[agent],
        tasks=[image_caption_task],
        verbose=True,
    )

    try:
        result = crew.kickoff()
        return result
    except Exception as e:
        return f"Error in processing the image: {str(e)}"

def generate_caption(image):
    try:
        unique_filename = f"{uuid.uuid4()}.png"
        image_path = os.path.join(SAVE_DIR, unique_filename)
        image.save(image_path)
        result = process_image_with_agent(image_path, image_caption_agent)
        return result
    except Exception as e:
        return f"Error: {str(e)}"

interface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Textbox(label="Generated Title and Description"),
    title="Image Captioning with LLM",
    description="Upload an image and get a title and description generated by an AI agent."

    
)

if __name__ == "__main__":
    interface.launch(share=True)