import os
import sys

import cv2
import numpy as np

import reconocedor as rec

CONOCIDOS = rec.CONOCIDOS
WEB_CAM = 0


def escribir(texto, frame):
    cv2.putText(frame, texto, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)


def capturar_cara(r, frame):
    caras = r.detectar_caras(frame)
    if not caras:
        return None, None
    mejor = max(caras, key=lambda c: c[-1])
    emb = r.embedding_cara(frame, mejor)
    return emb, mejor


def capturar_cuerpo(r, frame):
    personas = r._detectar_personas(frame)
    if not personas:
        return None, None
    mayor = max(personas, key=lambda p: p[2] * p[3])
    x, y, w, h = mayor
    recorte = frame[y:y + h, x:x + w]
    emb = r.embedding_cuerpo(recorte)
    return emb, mayor


def main():
    nombres = sys.argv[1:3]
    while len(nombres) < 2:
        nombres.append(f"Hija{len(nombres) + 1}")
    print(f"Registrando: {nombres}")

    r = rec.cargar_reconocedor()
    os.makedirs(CONOCIDOS, exist_ok=True)

    cap = cv2.VideoCapture(WEB_CAM)
    if not cap.isOpened():
        print("ERROR: no se pudo abrir la webcam")
        sys.exit(1)

    cv2.namedWindow("Registro", cv2.WINDOW_NORMAL)

    for nombre in nombres:
        carpeta = os.path.join(CONOCIDOS, nombre)
        os.makedirs(carpeta, exist_ok=True)
        print(f"\n>>> Hija: {nombre}")

        fase = "cara"
        while True:
            ok, frame = cap.read()
            if not ok:
                continue

            caras = r.detectar_caras(frame)
            personas = r._detectar_personas(frame)

            for c in caras:
                x, y, w, h = (int(v) for v in c[:4])
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            for p in personas:
                x, y, w, h = p
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)

            if fase == "cara":
                escribir("CARA: pon la cara mirando a la camara y pulsa ESPACIO", frame)
            else:
                escribir("CUERPO: alejate (cuerpo entero) y pulsa B", frame)
            escribir(f"Persona: {nombre}", frame)

            cv2.imshow("Registro", frame)
            tecla = cv2.waitKey(1) & 0xFF

            if tecla == ord("q"):
                cap.release()
                cv2.destroyAllWindows()
                print("Abortado por el usuario")
                sys.exit(0)

            if fase == "cara" and tecla == ord(" "):
                emb, caja = capturar_cara(r, frame)
                if emb is None:
                    print("  No se vio ninguna cara clara. Reintenta.")
                    continue
                np.save(os.path.join(carpeta, "face.npy"), emb)
                x, y, w, h = (int(v) for v in caja[:4])
                cv2.imwrite(os.path.join(carpeta, "cara_ref.jpg"),
                            frame[max(0, y - 20):y + h + 20, max(0, x - 20):x + w + 20])
                print(f"  Cara de {nombre} registrada (score {caja[-1]:.2f})")
                fase = "cuerpo"

            elif fase == "cuerpo" and tecla == ord("b"):
                emb, caja = capturar_cuerpo(r, frame)
                if emb is None:
                    print("  No se detecto ninguna persona. Reintenta.")
                    continue
                np.save(os.path.join(carpeta, "body.npy"), emb)
                x, y, w, h = caja
                cv2.imwrite(os.path.join(carpeta, "cuerpo_ref.jpg"),
                            frame[max(0, y):y + h, max(0, x):x + w])
                print(f"  Cuerpo de {nombre} registrado")
                break

    cap.release()
    cv2.destroyAllWindows()
    print("\nRegistro completado. Conocidos guardados:")
    for n in nombres:
        print(f"  - {n} (face.npy + body.npy)")


if __name__ == "__main__":
    main()