import os
import base64
import gradio as gr
from PIL import Image
import io
from groq import Groq

# Load environment variables
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

# Initialize Groq client
client = Groq(api_key=GROQ_API_KEY)

def analyze_construction_image(image):
    if image is None:
        return "Error: No image uploaded", "", ""

    try:
        # Convert PIL Image to base64
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        image_data_url = f"data:image/png;base64,{img_str}"

        # Prepare the message for Groq API
        completion = client.chat.completions.create(
            model="llama-3.2-11b-vision-preview",
            messages=[
                {
                    "role": "system",
                    "content": "You are an AI assistant specialized in analyzing construction site images. Identify issues, categorize them, and provide steps to resolve them."
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Analyze this construction image. Identify the snag category, provide a detailed snag description, and list steps to desnag."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": image_data_url
                            }
                        }
                    ]
                }
            ],
            temperature=0.7,
            max_tokens=300,
            top_p=1,
            stream=False,
            stop=None
        )

        result = completion.choices[0].message.content

        # Parse the result
        lines = result.split('\n')
        snag_category = lines[0] if len(lines) > 0 else "N/A"
        snag_description = lines[1] if len(lines) > 1 else "N/A"
        desnag_steps = "\n".join(lines[2:]) if len(lines) > 2 else "N/A"

        return snag_category, snag_description, desnag_steps
    except Exception as e:
        return f"Error: {str(e)}", "", ""

# Create the Gradio interface
iface = gr.Interface(
    fn=analyze_construction_image,
    inputs=gr.Image(type="pil", label="Upload Construction Image"),
    outputs=[
        gr.Textbox(label="Snag Category"),
        gr.Textbox(label="Snag Description"),
        gr.Textbox(label="Steps to Desnag")
    ],
    title="Construction Image Analyzer (Llama 3.2 Vision via Groq)",
    description="Upload a construction site image to identify issues and get desnag steps using Llama 3.2 Vision technology through Groq API.",
    examples=[
        ["example_image1.jpg"],
        ["example_image2.jpg"]
    ],
    cache_examples=True,
    theme="default"
)

# Launch the app
if __name__ == "__main__":
    iface.launch()