import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import requests
import torch
import os

# Load your Groq API key from Hugging Face Space secrets
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Load BLIP model for image captioning
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to("cuda" if torch.cuda.is_available() else "cpu")

# Function to generate image caption using BLIP
def generate_caption(image):
    inputs = processor(image, return_tensors="pt").to(model.device)
    output = model.generate(**inputs)
    caption = processor.decode(output[0], skip_special_tokens=True)
    return caption

# Function to get a repair suggestion from Groq API
def ask_groq(caption, question):
    prompt = f"""The following image caption describes some construction damage: "{caption}"

The user has a question about the damage: "{question}"

Based on the caption and the question, provide:
- Type of damage
- Tools needed
- Estimated repair time
- Step-by-step repair guidance

Be as detailed and practical as possible.
"""

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

    body = {
        "messages": [
            {"role": "system", "content": "You are a helpful civil engineer and construction repair assistant."},
            {"role": "user", "content": prompt}
        ],
        "model": "llama3-70b-8192"
    }

    response = requests.post("https://api.groq.com/openai/v1/chat/completions", headers=headers, json=body)

    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"Error: {response.status_code} - {response.text}"

# Main pipeline
def process_image_and_question(image, question):
    if image is None:
        return "Please upload an image."
    if not question:
        return "Please enter a question about the damage."

    caption = generate_caption(image)
    response = ask_groq(caption, question)
    return f"📸 Image Caption: {caption}\n\n💬 Response:\n{response}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🏗️ Construction Damage Assistant")
    gr.Markdown("Upload an image and ask a repair question. The AI will identify the issue and help you fix it.")

    with gr.Row():
        image = gr.Image(label="📷 Upload Damage Image", type="pil")
    question = gr.Textbox(label="💬 Your Question", placeholder="What tools do I need? How long will it take?")
    submit_btn = gr.Button("Get Repair Advice")
    output = gr.Textbox(label="🛠️ AI Response", lines=15)

    submit_btn.click(fn=process_image_and_question, inputs=[image, question], outputs=output)

# Launch the app
demo.launch()