import torch
import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration
import time
import logging

# Corrected the logging function name
logging.getLogger("asyncio").setLevel(logging.CRITICAL)

# --- CONFIGURATION ---
model_id = "Salesforce/blip-image-captioning-base"
device = "cuda" if torch.cuda.is_available() else "cpu"

# --- MODEL LOADING ---
print(f"🚀 TurboVision System starting on: {device}")
try:
    processor = BlipProcessor.from_pretrained(model_id)
    model = BlipForConditionalGeneration.from_pretrained(model_id).to(device)
    print("✅ SUCCESS: BLIP Model is ready!")
except Exception as e:
    print(f"❌ LOADING ERROR: {e}")
    model = None

# --- ANALYSIS LOGIC ---
def generate_knowledge(image):
    if model is None: 
        return "Model not loaded.", "N/A", "N/A", "0s"
    if image is None: 
        return "Please upload an image.", "N/A", "N/A", "0s"
    
    start_time = time.time()
    try:
        inputs = processor(image, return_tensors="pt").to(device)
        out = model.generate(**inputs)
        summary = processor.decode(out[0], skip_special_tokens=True)
        
        inputs_det = processor(image, "a detailed description of", return_tensors="pt").to(device)
        out_det = model.generate(**inputs_det)
        details = processor.decode(out_det[0], skip_special_tokens=True)
        
        speed = f"{round(time.time() - start_time, 2)}s"
        return summary.capitalize(), details.capitalize(), "Visual Analysis Complete.", speed
    except Exception as e:
        return f"Error: {str(e)}", "N/A", "N/A", "Error"

# --- UI DESIGN ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div style="text-align: center; background: #1e293b; padding: 30px; border-radius: 15px; margin-bottom: 25px; color: white; border-bottom: 5px solid #3b82f6;">
        <h1 style="font-size: 2.5em; margin: 0;">🚀 TurboVision AI</h1>
        <p style="color: #38bdf8; font-size: 1.2em; margin-top: 10px;">Deep Learning Image Intelligence by Muhammad Bilal</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(type="pil", label="Source Image")
            submit_btn = gr.Button("⚡ Run Intelligence Scan", variant="primary")
            gr.Markdown(f"**Backend:** Stable-BLIP Architecture  \n**Hardware:** {device.upper()} Mode")

        with gr.Column(scale=2):
            speed_lab = gr.Label(label="Processing Speed")
            out1 = gr.Textbox(label="Turbo Summary", placeholder="Summary will appear here...")
            out2 = gr.Textbox(label="Detailed Breakdown", placeholder="Object details will appear here...")
            out3 = gr.Textbox(label="System Status", value="Ready to Scan")

    submit_btn.click(
        fn=generate_knowledge, 
        inputs=input_img, 
        outputs=[out1, out2, out3, speed_lab]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)