"""Reconhecimento de pessoas: quem é quem, só na placa.

Cada foto de evento (o celular manda uma quando o sensor vê alguém) passa pelo YuNet, que acha os
rostos, e pelo SFace, que vira cada rosto numa "impressão" de 128 números. A impressão é comparada
com as das pessoas cadastradas (semelhança de cosseno); acima do limiar, é aquela pessoa. Rostos sem
nome ficam numa lista para alguém marcar no painel ("esta é a Ana"), e cada marcação vira amostra.

Privacidade: nada sai da placa. Fotos, recortes e impressões ficam em /opt/carinha-pc; apagar uma
pessoa apaga as amostras e os recortes dela; rostos sem nome somem depois de DIAS_SEM_NOME dias.
"""
import json
import os
import queue
import sqlite3
import threading
import time
from contextlib import contextmanager
from pathlib import Path

from comum import gravar_atomico, ler_json

LIMIAR = 0.40            # cosseno mínimo para dizer "é essa pessoa" (o SFace recomenda 0,363)
LIMIAR_GRUPO = 0.45      # rostos sem nome tão parecidos assim vão para o mesmo grupo (marcar de uma vez)
SCORE_ROSTO = 0.80       # confiança mínima do detector para considerar um rosto
TAMANHO_MIN = 48         # px: rosto menor que isso é longe demais para reconhecer bem
DIAS_SEM_NOME = 7        # rostos sem nome guardados para marcar no painel
MAX_AMOSTRAS = 30        # por pessoa (as mais novas)
PADRAO = dict(ligado=True)
MODELOS = ('face_detection_yunet_2023mar.onnx', 'face_recognition_sface_2021dec.onnx')


def cosseno(a, b):
    na = sum(x * x for x in a) ** 0.5
    nb = sum(x * x for x in b) ** 0.5
    if not na or not nb:
        return 0.0
    return sum(x * y for x, y in zip(a, b)) / (na * nb)


class Visao:
    """YuNet + SFace (OpenCV). Carrega na primeira foto; sem OpenCV ou modelo, fica indisponível."""

    def __init__(self, pasta_modelos):
        self.pasta = Path(pasta_modelos)
        self.det = self.rec = self.cv2 = self.np = None
        self.motivo = ''

    def pronta(self):
        if self.rec:
            return True
        try:
            import cv2
            import numpy as np
        except ImportError:
            self.motivo = 'OpenCV não instalado na placa (apt install python3-opencv)'
            return False
        faltando = [m for m in MODELOS if not (self.pasta / m).exists()]
        if faltando:
            self.motivo = 'Faltam os modelos em %s: %s' % (self.pasta, ', '.join(faltando))
            return False
        self.cv2, self.np = cv2, np
        self.det = cv2.FaceDetectorYN.create(str(self.pasta / MODELOS[0]), '', (320, 320), SCORE_ROSTO, 0.3, 5000)
        self.rec = cv2.FaceRecognizerSF.create(str(self.pasta / MODELOS[1]), '')
        self.motivo = ''
        return True

    def rostos(self, jpeg):
        """bytes JPEG -> [(vetor[128], recorte_jpeg, score, largura)]"""
        cv2, np = self.cv2, self.np
        if not jpeg:                                        # foto que não chegou inteira
            return []
        img = cv2.imdecode(np.frombuffer(jpeg, np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            return []
        h, w = img.shape[:2]
        escala = min(1.0, 640 / max(h, w))                  # a placa é pequena: analisa em até 640 px
        if escala < 1:
            img = cv2.resize(img, (int(w * escala), int(h * escala)))
            h, w = img.shape[:2]
        self.det.setInputSize((w, h))
        _, achados = self.det.detect(img)
        saida = []
        for f in (achados if achados is not None else []):
            if f[14] < SCORE_ROSTO or min(f[2], f[3]) < TAMANHO_MIN * escala:
                continue
            alinhado = self.rec.alignCrop(img, f)
            vetor = self.rec.feature(alinhado).flatten().tolist()
            ok, recorte = cv2.imencode('.jpg', alinhado, [cv2.IMWRITE_JPEG_QUALITY, 85])
            saida.append((vetor, recorte.tobytes() if ok else b'', float(f[14]), float(f[2])))
        return saida


class Reconhecimento:
    def __init__(self, pasta, visao=None, iniciar=True):
        self.pasta = Path(pasta)
        self.recortes = self.pasta / 'fotos' / 'rostos'
        self.recortes.mkdir(parents=True, exist_ok=True)
        self.caminho_db = str(self.pasta / 'pessoas.db')
        self.arquivo_config = self.pasta / 'pessoas.json'
        self.visao = visao or Visao(self.pasta / 'modelos')
        self.lock = threading.RLock()
        self.fila = queue.Queue(maxsize=500)
        self.publicar = lambda tipo, dados: None
        self.console = lambda origem, texto, motivo='', nivel='info', agora=None: None
        self.ao_reconhecer = lambda pessoa, confianca, agora: None   # o modo vida cumprimenta pelo nome
        self.varrendo = None                                         # {feitas, total} ao procurar nas fotos guardadas
        self.config = dict(PADRAO)
        try:
            self.config.update({k: v for k, v in json.loads(self.arquivo_config.read_text(encoding='utf-8')).items()
                                if k in PADRAO and isinstance(v, bool)})
        except (OSError, ValueError, AttributeError):
            pass
        with self.db() as db:
            db.execute('CREATE TABLE IF NOT EXISTS pessoas(id INTEGER PRIMARY KEY, nome TEXT UNIQUE, criado REAL, '
                       'visto REAL, vezes INTEGER DEFAULT 0, afinidade REAL DEFAULT 0)')
            db.execute('CREATE TABLE IF NOT EXISTS rostos(id INTEGER PRIMARY KEY, horario REAL, evento INTEGER, '
                       'pessoa INTEGER, amostra INTEGER DEFAULT 0, confianca REAL, score REAL, vetor TEXT, recorte TEXT)')
        if iniciar:
            threading.Thread(target=self.trabalhar, daemon=True).start()
            threading.Thread(target=self.faxina, daemon=True).start()

    def faxina(self):
        while True:
            time.sleep(3600)
            try:
                self.limpar_antigos()
            except Exception:
                pass

    @contextmanager
    def db(self):
        con = sqlite3.connect(self.caminho_db)
        try:
            with con:
                yield con
        finally:
            con.close()

    # ------------------------------------------------------------ pessoas
    def amostras(self):
        """{pessoa: [vetores]} das faces marcadas por alguém."""
        with self.db() as db:
            linhas = db.execute('SELECT pessoa, vetor FROM rostos WHERE amostra=1 AND pessoa IS NOT NULL').fetchall()
        saida = {}
        for pessoa, vetor in linhas:
            saida.setdefault(pessoa, []).append(json.loads(vetor))
        return saida

    def quem(self, vetor, amostras=None):
        """(pessoa, confiança) da amostra mais parecida, ou (None, melhor) abaixo do limiar."""
        melhor, quem = 0.0, None
        for pessoa, vetores in (amostras if amostras is not None else self.amostras()).items():
            s = max(cosseno(vetor, v) for v in vetores)
            if s > melhor:
                melhor, quem = s, pessoa
        return (quem, melhor) if melhor >= LIMIAR else (None, melhor)

    def pessoa(self, ident):
        with self.db() as db:
            linha = db.execute('SELECT id, nome, visto, vezes, afinidade FROM pessoas WHERE id=?', (ident,)).fetchone()
        return dict(id=linha[0], nome=linha[1], visto=linha[2], vezes=linha[3], afinidade=round(linha[4] or 0, 1)) if linha else None

    def pessoa_por_nome(self, nome):
        with self.db() as db:
            linha = db.execute('SELECT id FROM pessoas WHERE lower(nome)=lower(?)', (nome,)).fetchone()
        return self.pessoa(linha[0]) if linha else None

    # ------------------------------------------------------------ fotos
    def nova_foto(self, evento, caminho):
        """Chamado pelo histórico quando chega a foto de um evento: analisa em segundo plano."""
        if self.config['ligado']:
            try:
                self.fila.put_nowait((evento, str(caminho), False))
            except queue.Full:
                pass

    def analisar(self, evento, jpeg, agora=None, avisar=True):
        """Acha e reconhece os rostos de uma foto. Devolve [(pessoa ou None, confiança)]."""
        agora = agora or time.time()
        if not self.visao.pronta():
            return []
        achados = self.visao.rostos(jpeg)
        amostras = self.amostras()
        resultado = []
        for vetor, recorte, score, _ in achados:
            pessoa, conf = self.quem(vetor, amostras)
            antes = self.pessoa(pessoa)['visto'] if pessoa else None      # para a saudade de cada um
            with self.lock, self.db() as db:
                ident = db.execute('INSERT INTO rostos(horario, evento, pessoa, confianca, score, vetor) VALUES (?,?,?,?,?,?)',
                                   (agora, evento, pessoa, round(conf, 3), round(score, 3), json.dumps([round(x, 5) for x in vetor]))).lastrowid
                if recorte:
                    nome = 'rosto-%d.jpg' % ident
                    (self.recortes / nome).write_bytes(recorte)
                    db.execute('UPDATE rostos SET recorte=? WHERE id=?', (nome, ident))
                if pessoa:
                    db.execute('UPDATE pessoas SET visto=?, vezes=vezes+1 WHERE id=?', (agora, pessoa))
            resultado.append((pessoa, conf))
            if avisar:
                if pessoa:
                    p = dict(self.pessoa(pessoa), visto_antes=antes)
                    self.console('sentidos', 'Reconheceu %s' % p['nome'], 'semelhança %d%% (evento %s)' % (round(conf * 100), evento),
                                 agora=agora)
                    try:
                        self.ao_reconhecer(p, conf, agora)
                    except Exception:
                        pass
                else:
                    self.console('sentidos', 'Viu um rosto que ainda não conhece', 'parecido %d%% com alguém (precisa %d%%)' % (
                        round(conf * 100), round(LIMIAR * 100)), 'detalhe', agora)
        if achados and avisar:
            self.avisar()
        return resultado

    def trabalhar(self):
        while True:
            evento, caminho, varredura = self.fila.get()
            try:
                jpeg = Path(caminho).read_bytes()
                self.analisar(evento, jpeg, avisar=not varredura)
            except Exception as e:
                print('reconhecimento: erro:', e, flush=True)
            finally:
                if varredura and self.varrendo:
                    self.varrendo['feitas'] += 1
                    if self.varrendo['feitas'] >= self.varrendo['total']:
                        self.varrendo = None
                        self.console('sentidos', 'Terminou de procurar rostos nas fotos guardadas')
                        self.avisar()

    def varrer(self, eventos):
        """Procura rostos nas fotos já guardadas (para cadastrar pessoas logo de cara)."""
        if not self.visao.pronta():
            raise ValueError(self.visao.motivo)
        if self.varrendo:
            raise ValueError('Já estou procurando nas fotos guardadas')
        with self.db() as db:
            feitos = {r[0] for r in db.execute('SELECT DISTINCT evento FROM rostos')}
        alvos = [(e, c) for e, c in eventos if e not in feitos and c]
        if not alvos:
            raise ValueError('Nenhuma foto nova para analisar')
        self.varrendo = dict(feitas=0, total=len(alvos))
        for evento, caminho in alvos:
            self.fila.put((evento, str(caminho), True))
        self.avisar()
        return self.status()

    # ------------------------------------------------------------ marcar no painel
    def marcar(self, rostos, pessoa=None, nome=None):
        """Diz de quem são estes rostos (vira amostra). nome cria a pessoa; sem pessoa e nome, descarta."""
        if not isinstance(rostos, list) or not rostos or not all(type(r) is int for r in rostos):
            raise ValueError('Escolha pelo menos um rosto')
        with self.lock, self.db() as db:
            if nome is not None:
                nome = ' '.join(str(nome).split())[:30]
                if not nome:
                    raise ValueError('Digite o nome')
                existente = db.execute('SELECT id FROM pessoas WHERE lower(nome)=lower(?)', (nome,)).fetchone()
                pessoa = existente[0] if existente else db.execute(
                    'INSERT INTO pessoas(nome, criado) VALUES (?,?)', (nome, time.time())).lastrowid
            elif pessoa is not None and not db.execute('SELECT 1 FROM pessoas WHERE id=?', (pessoa,)).fetchone():
                raise ValueError('Pessoa não encontrada')
            marcas = ','.join('?' * len(rostos))
            if pessoa is None:                     # não é ninguém (ou foto ruim): some da lista
                for (arquivo,) in db.execute('SELECT recorte FROM rostos WHERE id IN (%s)' % marcas, rostos).fetchall():
                    self.apagar_recorte(arquivo)
                db.execute('DELETE FROM rostos WHERE id IN (%s)' % marcas, rostos)
            else:
                db.execute('UPDATE rostos SET pessoa=?, amostra=1 WHERE id IN (%s)' % marcas, [pessoa] + rostos)
                # guarda só as amostras mais novas de cada pessoa
                velhas = db.execute('SELECT id, recorte FROM rostos WHERE pessoa=? AND amostra=1 ORDER BY id DESC LIMIT -1 OFFSET ?',
                                    (pessoa, MAX_AMOSTRAS)).fetchall()
                for vid, arquivo in velhas:
                    self.apagar_recorte(arquivo)
                    db.execute('DELETE FROM rostos WHERE id=?', (vid,))
        if pessoa is not None:
            self.reconferir()
            p = self.pessoa(pessoa)
            self.console('sentidos', 'Aprendeu o rosto de %s' % p['nome'], '%d foto(s) marcada(s) no painel' % len(rostos))
        self.avisar()
        return self.status()

    def reconferir(self):
        """Com amostras novas, os rostos sem nome são comparados de novo (sem reabrir as fotos)."""
        amostras = self.amostras()
        if not amostras:
            return
        with self.lock, self.db() as db:
            for ident, vetor in db.execute('SELECT id, vetor FROM rostos WHERE pessoa IS NULL').fetchall():
                pessoa, conf = self.quem(json.loads(vetor), amostras)
                if pessoa:
                    db.execute('UPDATE rostos SET pessoa=?, confianca=? WHERE id=?', (pessoa, round(conf, 3), ident))

    def renomear(self, pessoa, nome):
        nome = ' '.join(str(nome or '').split())[:30]
        if not nome:
            raise ValueError('Digite o nome')
        with self.db() as db:
            try:
                db.execute('UPDATE pessoas SET nome=? WHERE id=?', (nome, pessoa))
            except sqlite3.IntegrityError:
                raise ValueError('Já existe alguém com esse nome') from None
        self.avisar()
        return self.status()

    def apagar_pessoa(self, pessoa):
        with self.lock, self.db() as db:
            for (arquivo,) in db.execute('SELECT recorte FROM rostos WHERE pessoa=?', (pessoa,)).fetchall():
                self.apagar_recorte(arquivo)
            db.execute('DELETE FROM rostos WHERE pessoa=?', (pessoa,))
            db.execute('DELETE FROM pessoas WHERE id=?', (pessoa,))
        self.avisar()
        return self.status()

    def gostou(self, pessoa, quanto=1.0):
        """Carinho ou conversa de quem foi reconhecido há pouco: a afinidade com essa pessoa cresce."""
        with self.db() as db:
            db.execute('UPDATE pessoas SET afinidade=min(100, afinidade+?) WHERE id=?', (quanto, pessoa))

    def apagar_recorte(self, arquivo):
        if arquivo:
            (self.recortes / arquivo).unlink(missing_ok=True)

    def limpar_antigos(self, agora=None):
        limite = (agora or time.time()) - DIAS_SEM_NOME * 86400
        with self.lock, self.db() as db:
            velhos = db.execute('SELECT id, recorte FROM rostos WHERE amostra=0 AND horario < ?', (limite,)).fetchall()
            for vid, arquivo in velhos:
                self.apagar_recorte(arquivo)
            db.execute('DELETE FROM rostos WHERE amostra=0 AND horario < ?', (limite,))
        return len(velhos)

    def configurar(self, dados):
        if not isinstance(dados, dict) or set(dados) - set(PADRAO) or not all(isinstance(v, bool) for v in dados.values()):
            raise ValueError('Ajuste inválido')
        self.config.update(dados)
        gravar_atomico(self.arquivo_config, self.config)
        self.avisar()
        return self.status()

    # ------------------------------------------------------------ painel
    def status(self):
        pronta = self.visao.pronta()
        with self.db() as db:
            pessoas = []
            for pid, nome, visto, vezes, afinidade in db.execute('SELECT id, nome, visto, vezes, afinidade FROM pessoas ORDER BY nome'):
                fotos = [dict(id=i, recorte=r) for i, r in db.execute(
                    'SELECT id, recorte FROM rostos WHERE pessoa=? AND recorte IS NOT NULL ORDER BY amostra DESC, id DESC LIMIT 6', (pid,))]
                amostras = db.execute('SELECT count(*) FROM rostos WHERE pessoa=? AND amostra=1', (pid,)).fetchone()[0]
                pessoas.append(dict(id=pid, nome=nome, visto=visto, vezes=vezes, afinidade=round(afinidade or 0, 1),
                                    amostras=amostras, fotos=fotos))
            linhas = db.execute('SELECT id, horario, evento, recorte, confianca, vetor FROM rostos WHERE pessoa IS NULL '
                                'AND recorte IS NOT NULL ORDER BY id DESC LIMIT 120').fetchall()
            total = db.execute('SELECT count(*) FROM rostos').fetchone()[0]
        sem_nome, grupos = [], []                 # grupos: rostos parecidos entre si (provável mesma pessoa)
        for i, hr, e, r, c, vetor in linhas:
            v = json.loads(vetor)
            g = next((g for g in grupos if cosseno(v, g['v']) >= LIMIAR_GRUPO), None)
            if g is None:
                g = dict(v=v, ids=[])
                grupos.append(g)
            g['ids'].append(i)
            sem_nome.append(dict(id=i, horario=hr, evento=e, recorte=r, parecido=c, grupo=grupos.index(g)))
        grupos = [g['ids'] for g in grupos]
        return dict(config=dict(self.config), disponivel=pronta, motivo='' if pronta else self.visao.motivo,
                    pessoas=pessoas, sem_nome=sem_nome, grupos=grupos, rostos=total, varrendo=self.varrendo, limiar=LIMIAR,
                    dias_sem_nome=DIAS_SEM_NOME)

    def avisar(self):
        self.publicar('pessoas', self.status())

    def rotas(self, h, metodo):
        rota = h.path.split('?', 1)[0]
        if not rota.startswith('/pessoas'):
            return False
        try:
            if metodo == 'GET' and rota == '/pessoas':
                resultado = self.status()
            elif metodo == 'GET' and rota.startswith('/pessoas/rosto/'):
                nome = os.path.basename(rota)
                arquivo = self.recortes / nome
                if not nome.startswith('rosto-') or not nome.endswith('.jpg') or not arquivo.is_file():
                    raise ValueError('Rosto não encontrado')
                h.enviar(arquivo.read_bytes(), 'image/jpeg')
                return True
            elif metodo == 'POST':
                d = ler_json(h, 8192)
                if rota == '/pessoas/marcar':
                    resultado = self.marcar(d.get('rostos'), d.get('pessoa'), d.get('nome'))
                elif rota == '/pessoas/renomear':
                    resultado = self.renomear(d.get('pessoa'), d.get('nome'))
                elif rota == '/pessoas/apagar':
                    resultado = self.apagar_pessoa(d.get('pessoa'))
                elif rota == '/pessoas/config':
                    resultado = self.configurar(d)
                elif rota == '/pessoas/varrer':
                    resultado = self.varrer(self.fotos_guardadas())
                else:
                    raise ValueError('Rota inválida')
            else:
                raise ValueError('Método inválido')
            h.enviar(json.dumps(resultado, ensure_ascii=False))
        except (ValueError, TypeError) as e:
            h.enviar(json.dumps(dict(erro=str(e)), ensure_ascii=False), codigo=400)
        return True

    fotos_guardadas = staticmethod(lambda: [])                       # o main liga ao histórico