import os
import gradio as gr
import requests
import base64
from PIL import Image
from io import BytesIO

# Set your Groq API key
GROQ_API_KEY = os.getenv("GROQ_API_KEY") or "PASTE_YOUR_GROQ_API_KEY_HERE"
GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"
GROQ_MODEL = "llama3-70b-8192"

def analyze_damage(image, user_prompt):
    if image is None or not user_prompt.strip():
        return "⚠️ Please upload an image and enter a question or description."

    # Convert image to base64 for potential future use
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    img_base64 = base64.b64encode(buffered.getvalue()).decode()

    system_prompt = (
        "You are a helpful construction engineer. The user has uploaded an image of some construction damage "
        "and has asked a question. Based on the image and their description or question, identify:\n"
        "- Type of damage\n"
        "- Possible causes\n"
        "- Recommended repair materials/tools\n"
        "- Estimated repair time"
    )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Image uploaded.\nQuestion: {user_prompt}"}
    ]

    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }

    payload = {
        "model": GROQ_MODEL,
        "messages": messages,
        "temperature": 0.7
    }

    try:
        response = requests.post(GROQ_URL, headers=headers, json=payload)
        response.raise_for_status()
        reply = response.json()["choices"][0]["message"]["content"]
        return reply
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as demo:
    gr.Markdown("# 🏗️ Construction Damage Assistant")

    gr.Markdown("""
    Upload an image of the damage and ask a question like:<br>
    🔹 *What type of crack is this?*<br>
    🔹 *What material should I use to fix it?*<br><br>
    The assistant will analyze your question and give expert suggestions.
    """, elem_classes="description-box")

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="📸 Upload Damage Photo", type="pil", height=280)

            question_input = gr.Textbox(
                label="🧾 Ask a Question or Describe the Problem",
                placeholder="e.g., What caused this crack? How do I fix it?",
                lines=3
            )

            submit_btn = gr.Button("🔍 Analyze", variant="primary")
            clear_btn = gr.Button("🧹 Clear Inputs")

        with gr.Column(scale=1.5):
            output = gr.Textbox(
                label="💡 AI Response",
                lines=18,
                show_copy_button=True
            )

    submit_btn.click(fn=analyze_damage, inputs=[image_input, question_input], outputs=output)
    clear_btn.click(lambda: (None, "", ""), outputs=[image_input, question_input, output])

    gr.Markdown("___")
    gr.Markdown("🔧 Powered by **Groq LLaMA3** | Built with ❤️ using **Gradio**")

demo.launch()