import gradio as gr
import cv2
import numpy as np
import os
import time
import gc
from inspector_engine import AdvancedBlockInspector

# Initialize engine with lazy loading
# Note: HF Spaces will run this on startup. 
# We use the local model file provided in the repository.
inspector = AdvancedBlockInspector(yolo_model_path='yolo26n-obb.pt')

def inspect(image):
    """Main inspection function"""
    if image is None:
        return None, {"error": "No image uploaded"}
    
    try:
        start_time = time.time()
        
        # Convert Gradio (RGB) to OpenCV (BGR)
        frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        
        # Process image
        result = inspector.inspect_block(frame)
        
        # Visualization
        vis_frame = frame.copy()
        if hasattr(inspector, 'last_saddles') and result.saddle_results:
            vis_frame = inspector.visualize_results(
                frame, 
                inspector.last_saddles, 
                result.saddle_results
            )
        
        # Convert back to RGB for Gradio
        vis_rgb = cv2.cvtColor(vis_frame, cv2.COLOR_BGR2RGB)
        
        # Prepare JSON data
        res_dict = result.to_dict()
        res_dict['server_side_time_ms'] = (time.time() - start_time) * 1000
        
        # Memory cleanup
        del frame, vis_frame
        gc.collect()
        
        return vis_rgb, res_dict
        
    except Exception as e:
        import traceback
        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
        print(error_msg)
        return None, {"error": str(e)}

# Create Gradio Interface with a premium theme
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo")) as demo:
    gr.Markdown("# 🔍 TMTL Industrial Inspector")
    gr.Markdown("### Remote AI Inference Engine for Saddle Defect Detection")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(type="numpy", label="Source Image")
            btn = gr.Button("🚀 Run Analysis", variant="primary")
        
        with gr.Column(scale=1):
            output_img = gr.Image(type="numpy", label="AI Visualization")
            output_json = gr.JSON(label="Detailed Analysis")
    
    gr.Markdown("---")
    gr.Markdown("© 2026 TMTL AI Solutions | Precision Inspection System")
    
    # Wire up the button with API name
    btn.click(
        fn=inspect, 
        inputs=input_img, 
        outputs=[output_img, output_json],
        api_name="predict"
    )

if __name__ == "__main__":
    demo.queue(
        max_size=10,
        default_concurrency_limit=4
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )