import sys, os
sys.path.append(os.path.dirname(__file__))

import gradio as gr
from PIL import Image
import pandas as pd

from inference.detect import detect_assets
from inference.classify import classify_qc

def run_detect(img, conf):
    if img is None:
        return None, "⚠️ Please upload an image.", None
        
    annotated, dets = detect_assets(img, conf=conf/100.0)
    
    if dets:
        top = dets[0]
        status = f"✅ **{top['name']}** — Confidence: {top['confidence']}%"
        df = pd.DataFrame(dets)
    else:
        status = "⚠️ No assets detected. Try another image."
        df = pd.DataFrame(columns=["name", "confidence", "box"])
        
    return annotated, status, df

def run_qc(img, product_type, threshold_pct):
    if img is None:
        return "", "⚠️ Please upload an image.", ""
        
    res = classify_qc(img, product_type=product_type, threshold=threshold_pct/100.0)
    
    warning_text = ""
    if res.get("warning"):
        warning_text = f"### ⚠️ {res['warning']}"
        
    status = res["status"]
    defect = res.get("defect", "")
    confidence = res["confidence"]
    source = res.get("source", "Unknown")
    
    if status == "QC Approved":
        result_text = f"### ✅ QC Status: **{status}** · Defect: **None** · Confidence: {confidence}%"
    else:
        result_text = f"### ❌ QC Status: **{status}** · Defect: **{defect}** · Confidence: {confidence}%"
        
    source_text = f"**ℹ️ Model Source:** {source}"
    
    return warning_text, result_text, source_text


with gr.Blocks(title="Yugensys Vision AI — Manufacturing") as demo:
    gr.Markdown("# Yugensys Vision AI")
    
    with gr.Tabs():
        # POC 1: Asset Identification
        with gr.TabItem("Asset Identification"):
            gr.Markdown("## Manufacturing Asset Identification")
            with gr.Row():
                with gr.Column(scale=1):
                    img_input1 = gr.Image(type="pil", label="Upload an image")
                    conf_slider = gr.Slider(minimum=0, maximum=100, value=25, step=1, label="Detection Confidence Threshold")
                    btn1 = gr.Button("Detect Assets", variant="primary")
                with gr.Column(scale=1):
                    img_output1 = gr.Image(label="Detections")
                    status_output1 = gr.Markdown()
                    table_output = gr.Dataframe(headers=["name", "confidence", "box"])
                    
            btn1.click(fn=run_detect, inputs=[img_input1, conf_slider], outputs=[img_output1, status_output1, table_output])
            
        # POC 2: Quality Inspection
        with gr.TabItem("Quality Inspection (Auto Fallback)"):
            gr.Markdown("## Visual Quality Inspection")
            with gr.Row():
                with gr.Column(scale=1):
                    img_input2 = gr.Image(type="pil", label="Upload a product image")
                    product_dropdown = gr.Dropdown(choices=["fmcg_inspection", "steel_inspection"], value="fmcg_inspection", label="Product Type")
                    thresh_slider = gr.Slider(minimum=0, maximum=100, value=80, step=1, label="EfficientNet Confidence Threshold (Triggers YOLO Fallback)")
                    btn2 = gr.Button("Run QC Inspection", variant="primary")
                with gr.Column(scale=1):
                    warning_output = gr.Markdown()
                    result_output = gr.Markdown()
                    source_output = gr.Markdown()
                    
            btn2.click(fn=run_qc, inputs=[img_input2, product_dropdown, thresh_slider], outputs=[warning_output, result_output, source_output])

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