import gradio as gr
import base64
import requests
import os

# ✅ LOAD API KEY from Space Secret
GROQ_API_KEY = os.getenv("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")
GROQ_MODEL = "llama3-70b-8192"
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"

def debug_env():
    if GROQ_API_KEY is None:
        return "❌ GROQ_API_KEY is NOT found. Did you set the Space Secret and restart the Space?"
    elif not GROQ_API_KEY.startswith("gsk_"):
        return f"⚠️ GROQ_API_KEY is loaded but seems malformed: {GROQ_API_KEY[:6]}..."
    else:
        return f"✅ GROQ_API_KEY loaded: {GROQ_API_KEY[:6]}..."

def image_to_base64(image_path):
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")
    except Exception:
        return None

def build_prompt(image_b64: str) -> str:
    return f"""
You are an expert in Non-Destructive Testing (NDT) and industrial defect analysis.
A user has uploaded an image of a defected component. Here's the base64 representation (partial): {image_b64[:300]}...
Analyze and return the following:
1. Defect Type
2. Recommended NDT Technique(s)
3. Explanation of Defect Cause
4. Solution/Fix
5. Time to Repair
6. Tools Required
7. Preventive Measures
"""

def query_groq(prompt: str) -> str:
    if not GROQ_API_KEY:
        return "❌ Error: GROQ_API_KEY is missing. Check Space secrets."
    if not GROQ_API_KEY.startswith("gsk_"):
        return "❌ Error: GROQ_API_KEY seems invalid. Should start with 'gsk_'."

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

    payload = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "system", "content": "You are a professional NDT inspector."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.5,
        "max_tokens": 800
    }

    try:
        response = requests.post(GROQ_API_URL, headers=headers, json=payload)
        if response.status_code == 401:
            return "❌ Error 401: Invalid API Key."
        elif response.status_code != 200:
            return f"❌ Error {response.status_code}: {response.text}"
        return response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        return f"❌ Exception: {str(e)}"

def analyze_defect(image_path):
    if not image_path:
        return "⚠️ Please upload an image."

    image_b64 = image_to_base64(image_path)
    if not image_b64:
        return "❌ Failed to read the image. Try again."

    prompt = build_prompt(image_b64)
    return query_groq(prompt)

# ==== GRADIO UI ====
with gr.Blocks() as demo:
    gr.Markdown("# 🔍 NDT Defect Analyzer")

    with gr.Tab("Analyze Image"):
        image_input = gr.Image(type="filepath", label="Upload Image")
        output_box = gr.Textbox(label="Analysis Output")
        analyze_btn = gr.Button("Analyze")
        analyze_btn.click(analyze_defect, inputs=image_input, outputs=output_box)

    with gr.Tab("Debug API Key"):
        debug_output = gr.Textbox(label="API Key Debug")
        debug_btn = gr.Button("Check")
        debug_btn.click(fn=debug_env, inputs=[], outputs=debug_output)

if __name__ == "__main__":
    demo.launch()