buffer.py

import os
import cv2
import shutil
import subprocess
import traceback   # <-- ADICIONE ISTO
from collections import deque
from datetime import datetime
from config import CAPTURE_DIR, FALLBACK_FPS

os.makedirs(CAPTURE_DIR, exist_ok=True)

class VideoBufferManager:
    """
    Circular pre-event buffer per trash_id and writes MP4 when an event fires.
    Agora com transcodificação automática para H.264 (web-friendly) via ffmpeg.
    """

    def __init__(self, buffer_seconds, post_seconds, frame_size, get_fps_fn,
                 on_video_closed=None, ffmpeg_bin="ffmpeg", ffprobe_bin="ffprobe",
                 transcode=True, replace_original=True):
        self.buffer_seconds   = float(buffer_seconds)
        self.post_seconds     = float(post_seconds)
        self.frame_size       = tuple(map(int, frame_size))  # (W, H)
        self.get_fps_fn       = get_fps_fn
        self.on_video_closed  = on_video_closed

        self.ffmpeg_bin       = ffmpeg_bin
        self.ffprobe_bin      = ffprobe_bin
        self.transcode        = bool(transcode)
        self.replace_original = bool(replace_original)

        # per-tid state
        self.buffers      = {}   # tid -> deque(frames)
        self.active       = {}   # tid -> bool
        self.post_count   = {}   # tid -> int (frames to write after event)
        self.open_writers = {}   # tid -> (cv2.VideoWriter, filepath)

    # ---- internals ---------------------------------------------------------

    def _ensure(self, tid, fps):
        fps = max(1.0, float(fps))
        maxlen = max(1, int(self.buffer_seconds * fps))
        if tid not in self.buffers:
            self.buffers[tid] = deque(maxlen=maxlen)
            self.active[tid] = False
            self.post_count[tid] = 0

    def _new_writer(self, tid, fps):
        # Ainda usamos 'mp4v' para capturar bruto (sem travar pipeline),
        # e depois transcodificamos para H.264 ao fechar o arquivo.
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        fname = os.path.join(CAPTURE_DIR, f"abandono_trash{tid}_{ts}.mp4")
        w, h = self.frame_size  # (W, H)
        writer = cv2.VideoWriter(fname, fourcc, float(fps), (w, h))
        if not writer.isOpened():
            raise RuntimeError(f"Failed to open VideoWriter: {fname} ({w}x{h}@{fps})")
        self.open_writers[tid] = (writer, fname)
        return writer, fname

    def _is_already_web_h264(self, path):
        """Retorna True se o vídeo já for H.264 (yuv420p)."""
        try:
            out = subprocess.check_output(
                [self.ffprobe_bin, "-v", "error",
                 "-select_streams", "v:0",
                 "-show_entries", "stream=codec_name,pix_fmt",
                 "-of", "default=noprint_wrappers=1:nokey=1",
                 path],
                stderr=subprocess.STDOUT
            ).decode("utf-8", "ignore").strip().splitlines()
            # Esperado: ["h264", "yuv420p"]
            if len(out) >= 2:
                codec, pix = out[0].strip(), out[1].strip()
                return codec == "h264" and pix.startswith("yuv420")
        except Exception:
            pass
        return False

    def _transcode_to_web(self, in_path, fps=None):
        """
        Transcodifica para H.264 (AVC) + yuv420p + faststart.
        Se replace_original=True, substitui o arquivo original de forma atômica.
        Retorna caminho do arquivo 'web' (final).
        """
        if not self.transcode:
            return in_path

        # Se já estiver em H.264/yuv420p, não precisa transcodificar
        if self._is_already_web_h264(in_path):
            return in_path

        base, ext = os.path.splitext(in_path)
        out_path = base + "-compat.mp4"

        # Monta comando ffmpeg
        cmd = [
            self.ffmpeg_bin, "-y",
            "-i", in_path,
            "-c:v", "libx264", "-profile:v", "high", "-level", "4.1",
            "-pix_fmt", "yuv420p",
            "-movflags", "+faststart",
            "-c:a", "aac", "-b:a", "128k",  # se não houver áudio, ffmpeg ignora
        ]
        if fps:
            cmd += ["-r", str(float(fps))]

        cmd += [out_path]

        try:
            subprocess.check_call(cmd)
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"ffmpeg failed for {in_path}: {e}") from e

        if self.replace_original:
            # Troca atômica: mantém o nome original apontando para a versão compatível
            tmp_backup = base + "-orig" + ext
            try:
                os.replace(in_path, tmp_backup)
                os.replace(out_path, in_path)
                # opcional: remover backup
                try:
                    os.remove(tmp_backup)
                except OSError:
                    pass
                return in_path
            except Exception as e:
                # se algo falhar, deixa o compat ao lado
                print(f"[WARN] Atomic replace failed: {e}. Using {out_path}")
                return out_path
        else:
            return out_path

    # ---- public API --------------------------------------------------------

    def maybe_start(self, tid, moving_pair):
        fps = self.get_fps_fn() or FALLBACK_FPS
        self._ensure(tid, fps)
        if moving_pair and not self.active.get(tid, False):
            self.active[tid] = True

    def push_frame(self, tid, frame):
        fps = self.get_fps_fn() or FALLBACK_FPS
        self._ensure(tid, fps)

        # Normalize frame size
        h, w = frame.shape[:2]
        fw, fh = self.frame_size
        if (w, h) != (fw, fh):
            frame = cv2.resize(frame, (fw, fh))

        # Buffer pre-event
        self.buffers[tid].append(frame.copy())

        # If writing, append and manage post count
        if tid in self.open_writers:
            writer, _ = self.open_writers[tid]
            writer.write(frame)

            if self.post_count.get(tid, 0) > 0:
                self.post_count[tid] -= 1
                if self.post_count[tid] == 0:
                    writer.release()
                    _, fname = self.open_writers[tid]
                    print(f"[CAPTURE] Closed: {fname}")
                    del self.open_writers[tid]
                    self.active[tid] = False
                    self.buffers[tid].clear()

                    # 🔁 Transcodificar para web
                    try:
                        web_path = self._transcode_to_web(fname, fps=fps)
                        print(f"[CAPTURE] Web-ready: {web_path}")
                        if callable(self.on_video_closed):
                            self.on_video_closed(tid, web_path)
                    except Exception as e:
                        import traceback
                        print(f"[WARN] Transcode failed for {fname}: {e}")
                        traceback.print_exc()
                        # Ainda assim, notificar com o original
                        if callable(self.on_video_closed):
                            self.on_video_closed(tid, fname)

    def on_event(self, tid):
        fps = self.get_fps_fn() or FALLBACK_FPS
        self._ensure(tid, fps)

        if tid not in self.open_writers:
            writer, fname = self._new_writer(tid, fps)
            # Dump pre-event history
            for fr in list(self.buffers.get(tid, [])):
                writer.write(fr)
            # Start post-event countdown (in frames)           
            MIN_POST_FRAMES = 30
            eff_fps = max(1.0, float(fps))  # evita zerar
            self.post_count[tid] = max(MIN_POST_FRAMES, int(self.post_seconds * eff_fps))
            print(f"[CAPTURE] Writing video: {fname}")

        self.active[tid] = True

    def close_all(self):
        # Fecha todos, transcodifica, e dispara callback
        for tid, (w, fname) in list(self.open_writers.items()):
            try:
                w.release()
                print(f"[CAPTURE] Closed: {fname}")
                fps = self.get_fps_fn() or FALLBACK_FPS
                try:
                    web_path = self._transcode_to_web(fname, fps=fps)
                    print(f"[CAPTURE] Web-ready: {web_path}")
                    if callable(self.on_video_closed):
                        self.on_video_closed(tid, web_path)
                except Exception as e:
                    import traceback
                    print(f"[WARN] Transcode failed for {fname}: {e}")
                    traceback.print_exc()
                    if callable(self.on_video_closed):
                        self.on_video_closed(tid, fname)
            except Exception as e:
                print(f"[WARN] Failed to close writer {fname}: {e}")
        self.open_writers.clear()

    def save_snapshot_clip(self, tid, fps=None, extra_post_frames=0):
        """
        Gera um clipe MP4 imediato com os frames atualmente no buffer do tid,
        transcodifica para web (H.264 + faststart) e retorna o caminho final.
        """
        fps = float(fps or self.get_fps_fn() or FALLBACK_FPS)
        frames = list(self.buffers.get(tid, []))
        if not frames:
            raise RuntimeError(f"No buffered frames for tid={tid} to snapshot.")

        if extra_post_frames > 0:
            last = frames[-1]
            frames.extend([last] * int(extra_post_frames))

        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        raw_name = os.path.join(CAPTURE_DIR, f"snapshot_trash{tid}_{ts}.mp4")
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        w, h = self.frame_size
        writer = cv2.VideoWriter(raw_name, fourcc, fps, (w, h))
        if not writer.isOpened():
            raise RuntimeError(f"Failed to open VideoWriter for snapshot: {raw_name} ({w}x{h}@{fps})")

        try:
            for fr in frames:
                fh, fw = fr.shape[:2]
                if (fw, fh) != (w, h):
                    fr = cv2.resize(fr, (w, h))
                writer.write(fr)
        finally:
            writer.release()

        print(f"[CAPTURE] Snapshot saved (raw): {raw_name}")

        # 🔁 Transcodificar para web e retornar caminho final
        final_path = self._transcode_to_web(raw_name, fps=fps)
        print(f"[CAPTURE] Snapshot web-ready: {final_path}")
        return final_path
    
    def push_frame_to_active_buffers(self, frame, active_ids: list):
        """
        Função otimizada para adicionar um único frame a múltiplos buffers ativos.
        Redimensiona o frame uma vez e o distribui.
        """
        if not active_ids:
            return

        fps = self.get_fps_fn() or FALLBACK_FPS

        # Otimização: Redimensiona o frame apenas uma vez
        h, w = frame.shape[:2]
        fw, fh = self.frame_size
        if (w, h) != (fw, fh):
            frame_resized = cv2.resize(frame, (fw, fh))
        else:
            frame_resized = frame

        # Distribui o frame para todos os buffers e writers ativos necessários
        for tid in active_ids:
            self._ensure(tid, fps)

            # Adiciona ao buffer de pré-evento
            self.buffers[tid].append(frame_resized.copy())

            # Se um vídeo para este ID já estiver sendo gravado, anexa o frame
            if tid in self.open_writers:
                writer, _ = self.open_writers[tid]
                writer.write(frame_resized)

                # Gerencia a contagem regressiva de frames pós-evento
                if self.post_count.get(tid, 0) > 0:
                    self.post_count[tid] -= 1
                    if self.post_count[tid] == 0:
                        writer.release()
                        _, fname = self.open_writers[tid]
                        print(f"[CAPTURE] Closed: {fname}")
                        del self.open_writers[tid]
                        self.active[tid] = False
                        self.buffers[tid].clear()

                        # Transcodifica para formato web e chama o callback
                        try:
                            web_path = self._transcode_to_web(fname, fps=fps)
                            print(f"[CAPTURE] Web-ready: {web_path}")
                            if callable(self.on_video_closed):
                                self.on_video_closed(tid, web_path)
                        except Exception as e:
                            print(f"[WARN] Transcode failed for {fname}: {e}")
                            traceback.print_exc()
                            if callable(self.on_video_closed):
                                self.on_video_closed(tid, fname)