#!/usr/bin/env python3
"""Rastreio das mãos no cut.mp4 (MediaPipe Hand Landmarker) -> assets/maos.json, para o "objeto preso na mão".

  uv run --python 3.12 --with mediapipe==0.10.21 --with numpy python ferramentas/rastrear_maos.py \
      --edit '<EDIT>' --de 8 --ate 13

Versão fixa de propósito: o mediapipe 1.x (e o Python 3.13) aborta com "Service is unavailable". A primeira vez
o uv baixa o Python 3.12 e o mediapipe (uns minutos, centenas de MB); avise a pessoa antes. Depois fica guardado.

--de/--ate limitam ao trecho que vai ter o objeto (mais rápido). O modelo hand_landmarker.task (Apache 2.0,
do Google) é procurado na skill do projeto (o "skill" do projeto.json, em assets/models), depois em
~/.editor-video/modelos (e na pasta antiga, ~/.sala-de-corte/modelos, sem baixar de novo); se não tiver, é baixado
na primeira vez para ~/.editor-video/modelos (uns 8 MB). Ou use --modelo.

Regras que já custaram retrabalho:
- Em cada corte do cut_timeline.json o rastreador recomeça: a mão de uma tomada não pode "pular" para a outra.
- O rótulo Left/Right do MediaPipe engana (depende de o vídeo estar espelhado). Aqui a mão é nomeada pelo lado
  da IMAGEM onde ela nasce e segue a trilha: "esquerda_da_imagem" ou "direita_da_imagem".
- A palma muitas vezes fica perto da borda de baixo. Por isso sai também "nos" (média dos nós dos dedos,
  pontos 5, 9, 13 e 17): o objeto fica acima deles.
- Suavização de 3 quadros centrada (atraso zero), sem atravessar corte.

Saída: {fps, largura, altura, cortes, quadros: [{t, maos: [{lado, palma, nos, pts}]}]}, coordenadas de 0 a 1.
"""
import argparse
import json
import os
import subprocess
import sys
import urllib.request

MODELO_URL = 'https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/latest/hand_landmarker.task'
AQUI = os.path.dirname(os.path.abspath(__file__))
PROJETO = os.path.dirname(AQUI)
PALMA = (0, 5, 9, 13, 17)


COMANDO = ("uv run --python 3.12 --with mediapipe==0.10.21 --with numpy python ferramentas/rastrear_maos.py "
           "--edit '<EDIT>' --de 8 --ate 13")
MEDIAPIPE = '0.10.21'


def skill_do_projeto():
    """(pasta da skill, de onde veio) pelo "skill" do projeto.json; senão a skill onde o modelo está; senão None"""
    try:
        with open(os.path.join(PROJETO, 'projeto.json'), encoding='utf-8') as f:
            s = str(json.load(f).get('skill') or '').strip()
        if s:
            return os.path.abspath(os.path.expanduser(s)), 'projeto.json ("skill")'
    except (OSError, ValueError, AttributeError):
        pass
    pai = os.path.dirname(PROJETO)
    if os.path.basename(PROJETO) == 'modelo-projeto' and os.path.exists(os.path.join(pai, 'helpers', 'matte.py')):
        return pai, 'a skill onde o modelo está'
    return None, ''


def modelo_padrao(skill=None):
    """a skill do projeto (assets/models), depois ~/.editor-video/modelos; o arquivo baixado antes na pasta antiga
    (~/.sala-de-corte/modelos) continua valendo"""
    if skill:
        na_skill = os.path.join(skill, 'assets', 'models', 'hand_landmarker.task')
        if os.path.exists(na_skill):
            return na_skill
    novo = os.path.join(os.path.expanduser('~'), '.editor-video', 'modelos', 'hand_landmarker.task')
    antigo = os.path.join(os.path.expanduser('~'), '.sala-de-corte', 'modelos', 'hand_landmarker.task')
    if not os.path.exists(novo) and os.path.exists(antigo):
        return antigo
    return novo
NOS = (5, 9, 13, 17)


def media(pts, idx):
    return [sum(pts[i][0] for i in idx) / len(idx), sum(pts[i][1] for i in idx) / len(idx)]


def main():
    ap = argparse.ArgumentParser(description='rastreio das mãos para prender um objeto na mão')
    ap.add_argument('--edit', required=True)
    ap.add_argument('--de', type=float, default=0.0)
    ap.add_argument('--ate', type=float, default=None)
    ap.add_argument('--modelo', default=None, help='hand_landmarker.task (padrão: a skill do projeto, ~/.editor-video/modelos)')
    ap.add_argument('--saida', default=os.path.join(PROJETO, 'assets', 'maos.json'))
    a = ap.parse_args()
    skill, origem = skill_do_projeto()
    print(f'skill: {skill} ({origem})' if skill else 'skill: nenhuma anotada no projeto.json ("skill"); o modelo de mãos '
          'vem de ~/.editor-video/modelos', flush=True)
    a.modelo = a.modelo or modelo_padrao(skill)
    try:
        import numpy as np
        import mediapipe as mp
        from mediapipe.tasks.python import BaseOptions, vision
    except ImportError:
        print(f'erro: falta o mediapipe. Rode assim (a primeira vez baixa uns minutos de pacotes):\n  {COMANDO}',
              file=sys.stderr)
        return 2
    versao = getattr(mp, '__version__', '?')
    if versao != MEDIAPIPE or sys.version_info[:2] != (3, 12):
        print(f'aviso: esta ferramenta foi conferida com o mediapipe {MEDIAPIPE} no Python 3.12 (aqui: mediapipe '
              f'{versao}, Python {sys.version_info[0]}.{sys.version_info[1]}). O mediapipe 1.x aborta com "Service is '
              f'unavailable". Se der erro, rode:\n  {COMANDO}', file=sys.stderr, flush=True)
    edit = os.path.abspath(os.path.expanduser(a.edit))
    with open(os.path.join(edit, 'cut_timeline.json'), encoding='utf-8') as f:
        tl = json.load(f)
    fps, W, H, dur = tl['fps'], tl['width'], tl['height'], tl['duration']
    ate = min(a.ate if a.ate is not None else dur, dur)
    f0, f1 = int(round(a.de * fps)), int(round(ate * fps))
    cortes = sorted(int(s['out_start_frame']) if 'out_start_frame' in s else int(round(s['out_start'] * fps))
                    for s in tl['segments'] if s['out_start'] > 0.01)
    if not os.path.exists(a.modelo):
        os.makedirs(os.path.dirname(a.modelo), exist_ok=True)
        print('baixando o modelo de mãos do MediaPipe (uns 8 MB)...', flush=True)
        urllib.request.urlretrieve(MODELO_URL, a.modelo)

    def novo():
        op = vision.HandLandmarkerOptions(base_options=BaseOptions(model_asset_path=a.modelo), num_hands=2,
                                          running_mode=vision.RunningMode.VIDEO, min_hand_detection_confidence=0.5,
                                          min_hand_presence_confidence=0.5, min_tracking_confidence=0.5)
        return vision.HandLandmarker.create_from_options(op)

    # quadros exatos, em sequência, por pipe (menor lado 540 basta para as mãos; as coordenadas saem de 0 a 1)
    esc = 540 / min(W, H)
    w2, h2 = int(round(W * esc / 2)) * 2, int(round(H * esc / 2)) * 2
    cmd = ['ffmpeg', '-v', 'error', '-ss', f'{f0 / fps:.4f}', '-i', os.path.join(edit, 'cut.mp4'), '-frames:v', str(f1 - f0),
           '-vf', f'scale={w2}:{h2}', '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-']
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    lm, cruas = novo(), {}
    n = f0
    while n < f1:
        buf = proc.stdout.read(w2 * h2 * 3)
        if len(buf) < w2 * h2 * 3:
            break
        if n in cortes and n != f0:
            lm.close()
            lm = novo()
        img = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.frombuffer(buf, np.uint8).reshape(h2, w2, 3))
        res = lm.detect_for_video(img, int(round(n * 1000 / fps)))
        cruas[n] = [[[p.x, p.y] for p in hand] for hand in res.hand_landmarks]
        n += 1
    proc.stdout.close()
    proc.wait()
    lm.close()
    nfim = n

    # trilhas: casa cada mão com a do quadro anterior pela palma, só dentro da mesma tomada
    trilhas, ativas, prox_id = {}, [], 0
    for n in range(f0, nfim):
        if n in cortes:
            ativas = []
        novas = []
        livres = list(ativas)
        for pts in cruas.get(n, []):
            pm = media(pts, PALMA)
            melhor = min(livres, key=lambda tr: (tr['ult'][0] - pm[0]) ** 2 + (tr['ult'][1] - pm[1]) ** 2, default=None)
            if melhor and ((melhor['ult'][0] - pm[0]) ** 2 + (melhor['ult'][1] - pm[1]) ** 2) ** 0.5 < 0.15 and n - melhor['n'] <= 3:
                livres.remove(melhor)
                tr = melhor
            else:
                tr = dict(id=prox_id, lado='esquerda_da_imagem' if pm[0] < 0.5 else 'direita_da_imagem', q={})
                prox_id += 1
            tr['ult'], tr['n'] = pm, n
            tr['q'][n] = pts
            trilhas[tr['id']] = tr
            novas.append(tr)
        ativas = novas + [tr for tr in livres if n - tr['n'] <= 3]
        lados = novas
        if len(lados) == 2 and lados[0]['lado'] == lados[1]['lado']:
            e, d = sorted(lados, key=lambda tr: tr['ult'][0])
            e['lado'], d['lado'] = 'esquerda_da_imagem', 'direita_da_imagem'
    quadros = {k: [] for k in range(f0, nfim)}
    for tr in trilhas.values():
        ks = sorted(tr['q'])
        for k in ks:   # média de 3 quadros centrada; na ponta da trilha fica o quadro cru
            viz = [tr['q'][j] for j in (k - 1, k, k + 1) if j in tr['q'] and not (j != k and (max(j, k) in cortes))]
            pts = [[sum(v[i][0] for v in viz) / len(viz), sum(v[i][1] for v in viz) / len(viz)] for i in range(21)]
            quadros[k].append(dict(lado=tr['lado'], palma=[round(x, 4) for x in media(pts, PALMA)],
                                   nos=[round(x, 4) for x in media(pts, NOS)], pts=[[round(x, 4), round(y, 4)] for x, y in pts]))
    saida = dict(fps=fps, largura=W, altura=H, cortes=[round(c / fps, 4) for c in cortes],
                 quadros=[dict(t=round(k / fps, 4), maos=quadros[k]) for k in sorted(quadros)])
    os.makedirs(os.path.dirname(os.path.abspath(a.saida)), exist_ok=True)
    with open(a.saida, 'w', encoding='utf-8') as f:
        json.dump(saida, f, separators=(',', ':'))
    com = sum(1 for k in quadros if quadros[k])
    print(f'ok: {a.saida} · {len(quadros)} quadros de {f0 / fps:.2f} a {nfim / fps:.2f} s · mão em {com} · {len(trilhas)} trilha(s)')
    return 0


if __name__ == '__main__':
    sys.exit(main())