import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

import cv2
import numpy as np
import time
import threading
from queue import Queue
from deepface import DeepFace
from ultralytics import YOLO
import warnings

warnings.filterwarnings("ignore")

# ──────────────────────────────────────────────────────────────────────────────
#  1. IMPORTAMOS NUESTROS MÓDULOS 
# ──────────────────────────────────────────────────────────────────────────────
# Del motor matemático y tracking
from seguimiento2 import GlobalMemory, CamManager, SECUENCIA, URLS, FUENTE

# Del motor de reconocimiento facial y audio
from reconocimiento2 import (
    gestionar_vectores,
    detectar_rostros_yunet,
    buscar_mejor_match,
    hilo_bienvenida,
    UMBRAL_SIM,
    COOLDOWN_TIME
)

# ──────────────────────────────────────────────────────────────────────────────
#  2. PROTECCIONES MULTIHILO E INICIALIZACIÓN
# ──────────────────────────────────────────────────────────────────────────────
COLA_ROSTROS = Queue(maxsize=4) 
YUNET_LOCK = threading.Lock() 
IA_LOCK = threading.Lock()    

# Inicializamos la base de datos usando tu función importada
print("\nIniciando carga de base de datos...")
BASE_DATOS_ROSTROS = gestionar_vectores(actualizar=True)

# ──────────────────────────────────────────────────────────────────────────────
#  3. MOTOR ASÍNCRONO
# ──────────────────────────────────────────────────────────────────────────────
def procesar_rostro_async(frame, box, gid, cam_id, global_mem, trk):
    """ Toma el recorte del tracker, usa YuNet importado, y hace la Fusión Mágica """
    try:
        if not BASE_DATOS_ROSTROS: return
        
        h_real, w_real = frame.shape[:2]
        escala_x = w_real / 480.0
        escala_y = h_real / 270.0
        
        x_min, y_min, x_max, y_max = box
        h_box = y_max - y_min
        
        y_min_expandido = max(0, y_min - (h_box * 0.15))
        y_max_cabeza = min(y_max, y_min + (h_box * 0.40)) 
        
        x1 = int(max(0, x_min) * escala_x)
        y1 = int(y_min_expandido * escala_y)
        x2 = int(min(480, x_max) * escala_x)
        y2 = int(y_max_cabeza * escala_y)
        
        roi_cabeza = frame[y1:y2, x1:x2]
        
        if roi_cabeza.size == 0 or roi_cabeza.shape[0] < 20 or roi_cabeza.shape[1] < 20: 
            return

        h_roi, w_roi = roi_cabeza.shape[:2]
        
        # Usamos la función de YuNet importada, pasándole nuestro candado
        faces = detectar_rostros_yunet(roi_cabeza, lock=YUNET_LOCK)

        # OJO: Tu detectar_rostros_yunet devuelve 5 valores (x, y, w, h, score)
        for (rx, ry, rw, rh, score) in faces:
            rx, ry = max(0, rx), max(0, ry)
            rw, rh = min(w_roi - rx, rw), min(h_roi - ry, rh)

            area_rostro_actual = rw * rh
            
            with global_mem.lock:
                data = global_mem.db.get(gid, {})
                nombre_actual = data.get('nombre')
                area_ref = data.get('area_rostro_ref', 0)
                
                necesita_saludo = False
                if str(cam_id) == "7": 
                    if not hasattr(global_mem, 'ultimos_saludos'):
                        global_mem.ultimos_saludos = {}
                    ultimo = global_mem.ultimos_saludos.get(nombre_actual if nombre_actual else "", 0)
                    if (time.time() - ultimo) > COOLDOWN_TIME:
                        necesita_saludo = True

            if nombre_actual is None or area_rostro_actual >= (area_ref * 1.5) or necesita_saludo:
                m_x = int(rw * 0.15)
                m_y = int(rh * 0.15)
                
                roi_rostro = roi_cabeza[max(0, ry-m_y):min(h_roi, ry+rh+m_y), 
                                        max(0, rx-m_x):min(w_roi, rx+rw+m_x)]
                
                if roi_rostro.size == 0 or roi_rostro.shape[0] < 40 or roi_rostro.shape[1] < 40: 
                    continue
                
                # ── Filtro de nitidez ──
                gray_roi = cv2.cvtColor(roi_rostro, cv2.COLOR_BGR2GRAY)
                nitidez = cv2.Laplacian(gray_roi, cv2.CV_64F).var()
                if nitidez < 50.0:
                    continue
                
                # ── ArcFace (Protegido con IA_LOCK) ──
                with IA_LOCK:
                    try:
                        res = DeepFace.represent(img_path=roi_rostro, model_name="ArcFace", enforce_detection=False)
                        emb = np.array(res[0]["embedding"], dtype=np.float32)
                        
                        # Usamos tu función importada (ya con producto punto)
                        mejor_match, max_sim = buscar_mejor_match(emb, BASE_DATOS_ROSTROS)
                    except Exception:
                        continue

                print(f"[DEBUG CAM {cam_id}] ArcFace: {mejor_match} al {max_sim:.2f} (Umbral: {UMBRAL_SIM})")

                if max_sim > UMBRAL_SIM and mejor_match:
                    nombre_limpio = mejor_match.split('_')[0]

                    with global_mem.lock:
                        gid_original = None
                        for otro_gid, datos_otro in global_mem.db.items():
                            if datos_otro.get('nombre') == nombre_limpio and otro_gid != gid:
                                gid_original = otro_gid
                                break
                        
                        if gid_original is not None:
                            print(f"\n[FUSIÓN MÁGICA] Uniendo el ID {gid} al original {gid_original} ({nombre_limpio})")
                            if gid in global_mem.db:
                                del global_mem.db[gid] 
                            
                            global_mem.db[gid_original]['ts'] = time.time()
                            global_mem.db[gid_original]['last_cam'] = cam_id
                            trk.gid = gid_original
                        else:
                            global_mem.db[gid]['nombre'] = nombre_limpio
                            global_mem.db[gid]['area_rostro_ref'] = area_rostro_actual

                        if str(cam_id) == "7" and necesita_saludo:
                            global_mem.ultimos_saludos[nombre_limpio] = time.time()
                            
                            try:
                                with IA_LOCK:
                                    analisis = DeepFace.analyze(roi_rostro, actions=['gender'], enforce_detection=False)[0]
                                    genero = analisis.get('dominant_gender', 'Man')
                            except Exception:
                                genero = "Man"
                            
                            # Usamos la función importada para el audio
                            threading.Thread(
                                target=hilo_bienvenida, 
                                args=(nombre_limpio, genero), 
                                daemon=True
                            ).start()
                    break 
    except Exception as e:
        pass
    finally:
        trk.procesando_rostro = False 

def worker_rostros(global_mem):
    """ Consumidor de la cola multihilo """
    while True:
        frame, box, gid, cam_id, trk = COLA_ROSTROS.get()
        procesar_rostro_async(frame, box, gid, cam_id, global_mem, trk)
        COLA_ROSTROS.task_done()

# ──────────────────────────────────────────────────────────────────────────────
#  4. LOOP PRINCIPAL DE FUSIÓN 
# ──────────────────────────────────────────────────────────────────────────────
class CamStream:
    def __init__(self, url):
        self.url = url
        self.cap = cv2.VideoCapture(url)
        self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        self.frame = None
        threading.Thread(target=self._run, daemon=True).start()
        
    def _run(self):
        while True:
            ret, f = self.cap.read()
            if ret: 
                self.frame = f
                time.sleep(0.01) 
            else: 
                time.sleep(2)
                self.cap.open(self.url)

def dibujar_track_fusion(frame_show, trk, global_mem):
    try: x1, y1, x2, y2 = map(int, trk.box)
    except Exception: return
    
    nombre_str = ""
    if trk.gid is not None:
        with global_mem.lock:
            nombre = global_mem.db.get(trk.gid, {}).get('nombre')
            if nombre: nombre_str = f" [{nombre}]"

    if trk.gid is None: color, label = (150, 150, 150), f"?{trk.local_id}"        
    elif nombre_str: color, label = (255, 0, 255), f"ID:{trk.gid}{nombre_str}" 
    elif trk.en_grupo: color, label = (0, 0, 255), f"ID:{trk.gid} [grp]"
    elif trk.aprendiendo: color, label = (255, 255, 0), f"ID:{trk.gid} [++]" 
    elif trk.origen_global: color, label = (0, 165, 255), f"ID:{trk.gid} [re-id]"
    else: color, label = (0, 255, 0), f"ID:{trk.gid}"
    
    cv2.rectangle(frame_show, (x1, y1), (x2, y2), color, 2)
    (tw, th), _ = cv2.getTextSize(label, FUENTE, 0.55, 1)
    cv2.rectangle(frame_show, (x1, y1-th-6), (x1+tw+2, y1), color, -1)
    cv2.putText(frame_show, label, (x1+1, y1-4), FUENTE, 0.55, (0,0,0), 1)

def main():
    print("\nIniciando Sistema")
    model = YOLO("yolov8n.pt")
    global_mem = GlobalMemory()
    managers = {str(c): CamManager(c, global_mem) for c in SECUENCIA}
    cams = [CamStream(u) for u in URLS]
    
    for _ in range(2):
        threading.Thread(target=worker_rostros, args=(global_mem,), daemon=True).start()

    cv2.namedWindow("SmartSoft", cv2.WINDOW_AUTOSIZE)
    idx = 0
    
    while True:
        now = time.time()
        tiles = []
        cam_ia = idx % len(cams)
        
        for i, cam_obj in enumerate(cams):
            frame = cam_obj.frame; cid = str(SECUENCIA[i])
            if frame is None: 
                tiles.append(np.zeros((270, 480, 3), np.uint8))
                continue
                
            frame_show = cv2.resize(frame.copy(), (480, 270))
            boxes = []
            turno_activo = (i == cam_ia)
            
            if turno_activo:
                res = model.predict(frame_show, conf=0.40, iou=0.50, classes=[0], verbose=False, imgsz=480)
                if res[0].boxes: 
                    boxes = res[0].boxes.xyxy.cpu().numpy().tolist()
                    
            tracks = managers[cid].update(boxes, frame_show, now, turno_activo)
            
            for trk in tracks:
                if trk.time_since_update <= 1: 
                    dibujar_track_fusion(frame_show, trk, global_mem)
                    
                    if turno_activo and trk.gid is not None and not getattr(trk, 'procesando_rostro', False):
                        if not COLA_ROSTROS.full():
                            trk.procesando_rostro = True
                            COLA_ROSTROS.put((frame.copy(), trk.box, trk.gid, cid, trk))

            if turno_activo: cv2.circle(frame_show, (460, 20), 6, (0, 0, 255), -1)
            
            con_id = sum(1 for t in tracks if t.gid and t.time_since_update==0)
            cv2.putText(frame_show, f"CAM {cid} [{con_id} ID]", (10, 28), FUENTE, 0.7, (255, 255, 255), 2)
            tiles.append(frame_show)
            
        if len(tiles) == 6: 
            cv2.imshow("SmartSoft Fusion", np.vstack([np.hstack(tiles[0:3]), np.hstack(tiles[3:6])]))
        
        idx += 1
        if cv2.waitKey(1) == ord('q'): 
            break
            
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()