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

# 🔍 Debug print (optional - remove later)
print("ENVIRONMENT VARIABLES:", list(os.environ.keys()))

# ✅ Correct way to load secret
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
MODEL_NAME = "llama-3.1-8b-instant"

# Convert image to base64 (not actually used in API, placeholder)
def image_to_base64(image: Image.Image) -> str:
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

# Main function to analyze defect
def analyze_defect(image, user_input):
    if image is None:
        return "⚠️ Please upload an image of the defective component."

    if not GROQ_API_KEY:
        return "❌ Environment variable 'GROQ_API_KEY' is set in the system but empty."

    try:
        # Optional: Keep for future image analysis
        image_b64 = image_to_base64(image)

        system_prompt = (
            "You are a mechanical inspection assistant trained in Non-Destructive Testing (NDT). "
            "You help identify defects in components such as pipes, welds, and metal surfaces. "
            "Given a user question and an image (image data not shown), choose the most appropriate "
            "NDT technique (X-ray, Ultrasonic, DPT, MPT), and suggest possible causes and next steps."
        )

        user_prompt = (
            f"The user uploaded an image of a mechanical component with visible defects. "
            f"The user says: '{user_input}'. Suggest the best NDT technique and possible solutions."
        )

        payload = {
            "model": MODEL_NAME,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            "temperature": 0.7
        }

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

        response = requests.post(GROQ_API_URL, headers=headers, json=payload)

        if response.status_code == 200:
            result = response.json()
            reply = result['choices'][0]['message']['content']
            return reply.strip()
        else:
            return f"❌ Groq API request failed. Status code: {response.status_code}"

    except Exception as e:
        return f"❌ An error occurred: {str(e)}"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## 🛠️ NDT Defect Analysis Assistant (Powered by Groq AI)")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")
        user_input = gr.Textbox(lines=4, label="Describe the issue or ask a question")

    analyze_button = gr.Button("Analyze Defect")
    output = gr.Textbox(label="AI's Response", lines=10)

    analyze_button.click(fn=analyze_defect, inputs=[image_input, user_input], outputs=output)

demo.launch()