import cv2
import torch
import time
import numpy as np
import mysql.connector
import face_recognition
from datetime import datetime

# --- CONFIGURAÇÃO ---
CAMERA_RTSP_URL = "rtsp://admin:senha123@192.168.1.64:554/stream1" # 👇 Altere para a sua câmera
CAMERA_ID_NO_DB = 1 # 👇 Coloque o ID desta câmera como está no seu DB `leitores`

DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'password': 'SUA_NOVA_SENHA_SEGURA', # 👈 Altere para a sua senha
    'database': 'facial'
}

# --- CARREGANDO OS MODELOS DE IA ---
print("Carregando modelo de detecção de pessoas (YOLOv5)...")
person_detector = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
person_detector.classes = [0]  # Apenas pessoas
person_detector.conf = 0.60    # Confiança mínima de 60%
print("Modelo de detecção de pessoas carregado.")

# --- BANCO DE DADOS E CACHE DE FACES CONHECIDAS ---
known_face_encodings = []
known_face_metadata = []

def get_db_connection():
    return mysql.connector.connect(**DB_CONFIG)

def carregar_faces_conhecidas():
    """ Carrega os vetores faciais da tabela 'pessoas' para a memória. """
    global known_face_encodings, known_face_metadata
    known_face_encodings.clear()
    known_face_metadata.clear()
    
    try:
        conn = get_db_connection()
        cursor = conn.cursor(dictionary=True)
        sql = """
            SELECT p.id as pessoa_id, p.nome_completo, f.vetor_facial
            FROM faces_cadastradas f
            JOIN pessoas p ON f.pessoa_id = p.id
            WHERE f.vetor_facial IS NOT NULL
        """
        cursor.execute(sql)
        for row in cursor.fetchall():
            encoding = np.frombuffer(row['vetor_facial'], dtype=np.float64)
            known_face_encodings.append(encoding)
            known_face_metadata.append({
                "pessoa_id": row['pessoa_id'],
                "nome_completo": row['nome_completo']
            })
        cursor.close()
        conn.close()
        print(f"Sucesso: {len(known_face_encodings)} faces conhecidas carregadas na memória.")
    except Exception as e:
        print(f"Erro ao carregar faces conhecidas do DB: {e}")

def registrar_alerta(pessoa_id, confianca):
    """ Insere um novo registro na tabela 'alertas_reconhecimento'. """
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        sql = """
            INSERT INTO alertas_reconhecimento 
            (leitor_id, pessoa_id, percentual_confianca, data_hora, imagem_capturada_url)
            VALUES (%s, %s, %s, %s, %s)
        """
        # A imagem_capturada_url pode ser implementada para salvar um snapshot do frame
        agora = datetime.now()
        valores = (CAMERA_ID_NO_DB, pessoa_id, confianca, agora, "placeholder.jpg")
        cursor.execute(sql, valores)
        conn.commit()
        cursor.close()
        conn.close()
        print(f"ALERTA REGISTRADO: Pessoa ID {pessoa_id} detectada com {confianca:.2f}% de confiança.")
    except Exception as e:
        print(f"Erro ao registrar alerta no DB: {e}")

# --- LOOP PRINCIPAL DE PROCESSAMENTO DE VÍDEO ---
carregar_faces_conhecidas()
last_alert_time = {} # Dicionário para evitar alertas repetidos da mesma pessoa

while True:
    try:
        print(f"Tentando conectar à câmera: {CAMERA_RTSP_URL}")
        cap = cv2.VideoCapture(CAMERA_RTSP_URL)
        if not cap.isOpened(): raise ValueError("Não foi possível abrir o stream.")
        print("Conexão bem-sucedida!")
        
        while True:
            ret, frame = cap.read()
            if not ret: break

            # 1. DETECÇÃO DE PESSOAS (YOLO)
            results = person_detector(frame)
            person_detections = results.xyxy[0]

            # 2. PARA CADA PESSOA, PROCURAR E RECONHECER O ROSTO
            for det in person_detections:
                x1, y1, x2, y2, conf, cls = det
                if conf >= person_detector.conf:
                    # Recorta a imagem da pessoa
                    pessoa_img = frame[int(y1):int(y2), int(x1):int(x2)]

                    # Encontra a localização do rosto dentro da imagem da pessoa
                    face_locations = face_recognition.face_locations(pessoa_img)
                    if face_locations:
                        # Gera o vetor (encoding) do primeiro rosto encontrado
                        face_encodings = face_recognition.face_encodings(pessoa_img, face_locations)
                        
                        if face_encodings:
                            unknown_encoding = face_encodings[0]
                            # Compara o rosto com todos os rostos conhecidos
                            matches = face_recognition.compare_faces(known_face_encodings, unknown_encoding, tolerance=0.6)
                            
                            nome_display = "Desconhecido"
                            cor_caixa = (0, 0, 255) # Vermelho para desconhecido
                            
                            if True in matches:
                                first_match_index = matches.index(True)
                                metadata = known_face_metadata[first_match_index]
                                nome_display = metadata['nome_completo']
                                cor_caixa = (0, 255, 255) # Amarelo para conhecido
                                
                                # Lógica para evitar alertas repetidos a cada segundo
                                pessoa_id = metadata['pessoa_id']
                                agora = time.time()
                                if pessoa_id not in last_alert_time or (agora - last_alert_time[pessoa_id]) > 60: # Cooldown de 60s
                                    face_distances = face_recognition.face_distance(known_face_encodings, unknown_encoding)
                                    confianca = (1 - face_distances[first_match_index]) * 100
                                    registrar_alerta(pessoa_id, confianca)
                                    last_alert_time[pessoa_id] = agora
                            
                            # Desenha o retângulo ao redor da pessoa e o nome
                            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), cor_caixa, 2)
                            cv2.putText(frame, nome_display, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, cor_caixa, 2)

            # Mostra o vídeo processado em uma janela
            cv2.imshow('FacialSys - Vigilância Ativa', frame)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                cap.release()
                cv2.destroyAllWindows()
                exit()
                
    except Exception as e:
        print(f"Ocorreu um erro no loop principal: {e}")
        print("Tentando reconectar em 10 segundos...")
        time.sleep(10)