# main.py
import os
import gradio as gr
from PIL import Image
import io
import base64
from groq import Groq

# Initialize Groq client with API key (set this as a secret in HF Spaces)
client = Groq(api_key=os.environ.get("construction"))

# Helper: Convert PIL Image to base64
def image_to_base64(image):
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode()

# Prompt for model
SYSTEM_PROMPT = """
You are a helpful civil engineering assistant. The user uploads an image showing some construction damage such as cracks, water leakage, or pipe failure. Based on the image, give:

1. Likely issue  
2. Possible solution  
3. Tools or materials needed  
4. Estimated time to fix

Use simple, helpful, practical language.
"""

# Chatbot logic
def analyze_image(image, history):
    if image is None:
        return history + [("User", "No image uploaded."), ("Bot", "Please upload a damage photo.")]

    base64_img = image_to_base64(image)
    image_url = f"data:image/jpeg;base64,{base64_img}"

    try:
        response = client.chat.completions.create(
            model="meta-llama/llama-4-scout-17b-16e-instruct",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": [
                    {"type": "text", "text": "Please analyze this image and give advice on the damage."},
                    {"type": "image_url", "image_url": {"url": image_url}}
                ]}
            ],
            temperature=0.7,
            max_tokens=512
        )

        reply = response.choices[0].message.content
        history.append(("User", "Uploaded image"))
        history.append(("Bot", reply))
        return history
    except Exception as e:
        return history + [("Bot", f"❌ Error: {str(e)}")]

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## 🛠️ Construction Damage Assistant\nUpload a photo of damage to get repair advice.")

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

        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Repair Suggestions", height=450)

    state = gr.State([])
    submit_btn = gr.Button("Analyze")

    submit_btn.click(fn=analyze_image, inputs=[image_input, state], outputs=chatbot)

demo.launch()