import streamlit as st
import cv2
import numpy as np
from PIL import Image
from yolo_inference import build_detector_from_env
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, RTCConfiguration
import av

# Configuração inicial da página do Streamlit (Título e Layout)
st.set_page_config(page_title="YOLO Detection - Streamlit", layout="wide", page_icon="🚀")

# Configuração RTC para STUN servers (mais robusta para evitar timeouts)
RTC_CONFIGURATION = RTCConfiguration(
    {
        "iceServers": [
            {"urls": ["stun:stun.l.google.com:19302"]},
            {"urls": ["stun:stun1.l.google.com:19302"]},
            {"urls": ["stun:stun2.l.google.com:19302"]},
            {"urls": ["stun:stun3.l.google.com:19302"]},
            {"urls": ["stun:stun4.l.google.com:19302"]},
            {"urls": ["stun:stun.services.mozilla.com"]},
        ]
    }
)

# Lista de classes do dataset personalizado para monitoramento especial
CUSTOM_CLASSES = {"car", "truck", "bus", "motorbike", "bicycle", "van", "threewheel"}

class YoloVideoProcessor(VideoProcessorBase):
    def __init__(self, detector):
        self.detector = detector

    def recv(self, frame):
        img = frame.to_ndarray(format="bgr24")
        
        # Realiza a detecção de objetos
        detections = self.detector.detect(img)
        
        # Desenha os resultados no frame
        img_out = self.detector.draw(img, detections)
        
        # Adiciona overlay de instrução
        cv2.putText(img_out, "YOLO Real-time Detection", (20, 40), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        return av.VideoFrame.from_ndarray(img_out, format="bgr24")

def main():
    """
    Função principal que gerencia a interface Streamlit.
    Permite alternar entre detecção em imagens estáticas e vídeo em tempo real via webcam.
    """
    st.title("🚀 YOLO Object Detection")
    st.markdown("---")
    st.markdown("### Interface interativa para detecção de objetos usando YOLOv3-tiny.")

    # Sidebar: Painel lateral para controle de parâmetros e seleção de modo
    st.sidebar.header("🛠️ Configurações do Modelo")
    
    # Sliders para ajuste dinâmico dos limiares de detecção
    conf_threshold = st.sidebar.slider("Confiança Mínima (Threshold)", 0.0, 1.0, 0.5, 0.05, 
                                     help="Nível mínimo de certeza para exibir uma detecção.")
    nms_threshold = st.sidebar.slider("NMS Threshold", 0.0, 1.0, 0.4, 0.05,
                                    help="Limiar para supressão de não-máximos (remove bboxes sobrepostas).")
    
    st.sidebar.markdown("---")
    # Seleção do modo de operação
    mode = st.sidebar.radio(
        "📡 Escolha o Modo de Entrada", 
        ["Imagem", "Câmera (Snapshot)", "Câmera (WebRTC)"],
        help="Snapshot: Mais estável em nuvem. WebRTC: Vídeo em tempo real (pode falhar em algumas redes)."
    )

    # Inicializa o detector YOLO
    try:
        detector = build_detector_from_env(conf_threshold=conf_threshold, nms_threshold=nms_threshold)
    except Exception as e:
        st.error(f"❌ Erro ao inicializar detector: {e}")
        return

    if mode == "Imagem":
        st.subheader("📁 Upload e Detecção em Imagem")
        uploaded_file = st.file_uploader("Arraste ou selecione uma imagem...", type=["jpg", "jpeg", "png"])

        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            process_and_display_image(image, detector)

    elif mode == "Câmera (Snapshot)":
        st.subheader("📸 Detecção via Foto (Snapshot)")
        st.info("Este modo é o mais estável para uso em nuvem. Ele tira uma foto e processa a detecção.")
        img_file = st.camera_input("Tirar Foto")
        
        if img_file:
            image = Image.open(img_file)
            process_and_display_image(image, detector)

    elif mode == "Câmera (WebRTC)":
        st.subheader("🎥 Detecção via Webcam (WebRTC)")
        st.info("Vídeo em tempo real. Se a conexão demorar (timeout), use o modo 'Câmera (Snapshot)'.")
        
        webrtc_streamer(
            key="yolo-detection",
            video_processor_factory=lambda: YoloVideoProcessor(detector),
            rtc_configuration=RTC_CONFIGURATION,
            media_stream_constraints={"video": True, "audio": False},
            async_processing=True,
            video_html_attrs={
                "style": {"width": "100%"},
                "controls": False,
                "autoPlay": True,
            },
        )

def process_and_display_image(image, detector):
    """Função auxiliar para processar e exibir os resultados de uma imagem PIL"""
    image_np = np.array(image)
    frame_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    with st.spinner('Processando...'):
        detections = detector.detect(frame_bgr)
    
    hits = sorted({d['class_name'] for d in detections if d['class_name'] in CUSTOM_CLASSES})
    
    col1, col2 = st.columns(2)
    with col1:
        st.image(image, caption="Entrada", use_column_width=True)
    with col2:
        result_bgr = detector.draw(frame_bgr, detections)
        result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
        st.image(result_rgb, caption="Resultado da Detecção", use_column_width=True)

    if hits:
        st.success(f"✅ Objetos detectados: **{', '.join(hits)}**")
    else:
        st.info("ℹ️ Nenhuma classe de interesse detectada.")

if __name__ == "__main__":
    main()