import os
import json
from typing import Optional
from datetime import datetime
import google.generativeai as genai
import gradio as gr
import PIL.Image
import tempfile
import requests
import base64
import io

# 🔑 Add your API key here
GEMINI_API_KEY = "AIzaSyDCLrgUo2RLpS0ShuoQFoLO00OqTgMVDs4"
MODEL_NAME = "models/gemini-2.5-flash"

# Configure API
genai.configure(api_key=GEMINI_API_KEY)


class GeminiChatBot:
    """Mechanical Component Defect Detection System"""

    def __init__(self, model_name: str = MODEL_NAME):
        self.model_name = model_name
        self.conversation_history = []
        self.system_prompt = ""

    def set_system_prompt(self):
        """Fixed technical prompt (no mode selection)"""
        self.system_prompt = """You are an expert mechanical inspection AI.

Analyze mechanical components for defects such as:
- Cracks
- Corrosion
- Wear and tear
- Misalignment
- Surface damage

Return precise engineering output only."""

    def chat(self, image: Optional[PIL.Image.Image] = None, temperature: float = 0.3) -> str:
        """Image-based defect detection"""

        try:
            self.set_system_prompt()

            if not image:
                return "Please upload an image."

            # Convert image → base64
            buffered = io.BytesIO()
            image.save(buffered, format="JPEG")
            image_data = base64.b64encode(buffered.getvalue()).decode("utf-8")

            # 🔥 DEFECT DETECTION PROMPT
            user_message = """
Analyze the given mechanical component image.

Task:
1. Determine whether the component is DEFECTIVE or NOT DEFECTIVE.
2. Detect cracks, corrosion, wear, deformation, misalignment.

Strict Output Format:
Status: <Defective / Not Defective>
Reason: <Short explanation>
Confidence: <0-100%>
"""

            contents = [
                {
                    "role": "user",
                    "parts": [
                        {
                            "inline_data": {
                                "mime_type": "image/jpeg",
                                "data": image_data
                            }
                        },
                        {
                            "text": f"[SYSTEM: {self.system_prompt}]\n\n{user_message}"
                        }
                    ]
                }
            ]

            url = f"https://generativelanguage.googleapis.com/v1beta/{MODEL_NAME}:generateContent?key={GEMINI_API_KEY}"

            payload = {
                "contents": contents,
                "generationConfig": {
                    "temperature": temperature,
                    "maxOutputTokens": 1000
                }
            }

            headers = {"Content-Type": "application/json"}

            response = requests.post(url, json=payload, headers=headers, timeout=30)
            response.raise_for_status()

            result = response.json()

            if "candidates" not in result:
                return "Error: No response from API"

            return result["candidates"][0]["content"]["parts"][0]["text"]

        except Exception as e:
            return f"Error: {str(e)}"


# Initialize chatbot
chatbot = GeminiChatBot()


# 🔁 Response function (NO TEXT INPUT)
def respond(image, chat_history, temperature):
    response = chatbot.chat(image=image, temperature=temperature)

    content = ""

    if image:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
            image.save(f.name)
            img_path = f.name.replace("\\", "/")
        content += f"![]({img_path})\n"

    chat_history.append({"role": "user", "content": content})
    chat_history.append({"role": "assistant", "content": response})

    return None, chat_history


def clear_history():
    chatbot.conversation_history = []
    return []


def export_chat(chat_history):
    if not chat_history:
        return "No data"

    return json.dumps({
        "timestamp": datetime.now().isoformat(),
        "conversation": chat_history
    }, indent=2)


# 🎨 Gradio UI (CLEAN VERSION)
with gr.Blocks(title="Mechanical Defect Detection", theme=gr.themes.Soft()) as demo:

    gr.Markdown("""
# 🔧 Mechanical Component Defect Detection System
Upload an image to detect defects using AI
""")

    chatbot_ui = gr.Chatbot(height=500)

    temperature = gr.Slider(0, 1, value=0.3, label="Temperature")

    img_input = gr.Image(
        type="pil",
        label="Upload Mechanical Component",
        sources=["upload", "clipboard"]
    )

    with gr.Row():
        analyze_btn = gr.Button("Analyze", variant="primary")
        clear_btn = gr.Button("Clear")
        export_btn = gr.Button("Export")

    export_output = gr.Textbox(visible=False)

    # 🔘 Actions
    analyze_btn.click(
        respond,
        inputs=[img_input, chatbot_ui, temperature],
        outputs=[img_input, chatbot_ui]
    )

    clear_btn.click(
        clear_history,
        outputs=[chatbot_ui]
    )

    export_btn.click(
        lambda x: (export_chat(x), gr.update(visible=True)),
        inputs=[chatbot_ui],
        outputs=[export_output, export_output]
    )


# Run app
if __name__ == "__main__":
    demo.launch()