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

# ✅ Load Groq API key from environment (add this in Hugging Face Secrets)
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_MODEL = "llama3-70b-8192"

# ✅ Load BLIP model for image captioning
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# 📸 Generate caption from image

def get_caption(image):
    inputs = processor(images=image, return_tensors="pt")
    out = blip_model.generate(**inputs)
    caption = processor.decode(out[0], skip_special_tokens=True)
    return caption

# 💬 Communicate with Groq API

def generate_response(image, user_query=""):
    if GROQ_API_KEY is None:
        return "❌ API key not set. Please set GROQ_API_KEY using Hugging Face Secrets."

    caption = get_caption(image)

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

    data = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "system", "content": "You are a helpful assistant that analyzes construction damage from image captions and gives expert repair advice."},
            {"role": "user", "content": f"Image caption: {caption}"},
            {"role": "user", "content": f"My question is: {user_query or 'No follow-up question.'}"},
            {"role": "user", "content": "Please provide:\n1. Type of damage\n2. Likely cause\n3. Suggested repair solutions\n4. Required tools\n5. Estimated time and cost"}
        ]
    }

    try:
        response = requests.post("https://api.groq.com/openai/v1/chat/completions", headers=headers, json=data, timeout=30)
        if response.status_code == 200:
            content = response.json()['choices'][0]['message']['content']
            return content if content.strip() else "⚠️ No response received. Please try again."
        else:
            return f"❌ Error {response.status_code}: {response.text}"
    except Exception as e:
        return f"❌ Exception: {str(e)}"

# ============================
# 🌟 Gradio Modern UI
# ============================

with gr.Blocks(css="""
.gradio-container {
  background-color: #f8f9fa;
  font-family: 'Segoe UI', sans-serif;
}
#title {
  text-align: center;
  font-size: 32px;
  font-weight: bold;
  padding: 20px 0;
  color: #333;
}
#subtitle {
  text-align: center;
  font-size: 16px;
  color: #666;
  margin-bottom: 20px;
}
.card {
  border: 1px solid #ddd;
  border-radius: 16px;
  background-color: white;
  padding: 20px;
  box-shadow: 0 4px 8px rgba(0,0,0,0.03);
}
""") as demo:

    gr.Markdown("<div id='title'>\ud83c\udfd7\ufe0f BuildFix AI — Construction Damage Inspector</div>")
    gr.Markdown("<div id='subtitle'>Upload a construction damage image and get expert analysis with repair suggestions instantly.</div>")

    with gr.Row():
        with gr.Column(scale=1, elem_classes="card"):
            image_input = gr.Image(type="pil", label="\ud83d\udcf8 Upload Image of Damage")
            user_query = gr.Textbox(label="\ud83d\udcac Ask a follow-up question (optional)", placeholder="e.g., How expensive is this repair?", lines=2)
            submit_btn = gr.Button("\ud83d\udd0d Analyze", variant="primary")

        with gr.Column(scale=1, elem_classes="card"):
            result_output = gr.Textbox(label="\ud83e\uddd0 AI Report", lines=20)

    submit_btn.click(fn=generate_response, inputs=[image_input, user_query], outputs=result_output)

demo.launch()