"""
Historical 2D Studio - Seven Nguyen Animated Tactical Actor (2026 Engine)
Lead Motion & Pipeline Automation Engineer: Elena Rostova
Head of Audio & Multilingual Localization: Arthur Pendelton

Features:
- Full Autonomic Animation System:
  + Natural Idle Breathing: sinusoidal displacement and subtle chest volume scaling
  + Natural Blinking Controller: 3-phase realistic eye lid closure every 3.5 - 4.5 seconds
  + Phonetic Lip-Sync: Viseme switching (ah, oh, closed, smile) via audio RMS or procedural cadence
  + Micro Head Nodding: Tone-aligned emphasis gestures (0 - 2px)
  + Seamless Skin-Matched Lip Patches with Alpha Feathering and Poisson Blending
- Supports all 5 production poses: p1 (welcome), p2 (pointing right), p3 (pointing left), p4 (crossed arms), p6 (leaning/thinking)
"""

import os
import math
import subprocess
import wave
import numpy as np
import cv2

def draw_natural_eye_blink(base_bgr, lx, ly, rx, ry, rw=6, rh=3, ratio=1.0):
    """
    Vẽ mi mắt chớp tự nhiên trực tiếp lên khuôn mặt:
    - lx, ly: tâm đồng tử mắt trái (trên ảnh)
    - rx, ry: tâm đồng tử mắt phải (trên ảnh)
    - rw, rh: bán kính hốc mắt (width/height)
    - ratio: tỷ lệ nhắm mắt (0.0 -> 1.0)
    """
    if ratio <= 0.02:
        return base_bgr
        
    res = base_bgr.copy()
    bh, bw = base_bgr.shape[:2]
    
    for cx, cy in [(int(lx), int(ly)), (int(rx), int(ry))]:
        if cx - rw < 0 or cx + rw >= bw or cy - rh < 0 or cy + rh >= bh:
            continue
            
        sample_y = min(bh - 1, max(0, cy + rh + 2))
        skin_color = base_bgr[sample_y, cx].astype(float)
        
        mask = np.zeros((bh, bw), dtype=np.uint8)
        cv2.ellipse(mask, (cx, cy), (rw, max(1, int(rh * ratio))), 0, 0, 360, 255, -1)
        
        alpha = (mask.astype(float) / 255.0) * ratio
        alpha_3d = np.repeat(alpha[:, :, None], 3, axis=2)
        res = (skin_color * alpha_3d + res.astype(float) * (1.0 - alpha_3d)).astype(np.uint8)
        
        if ratio > 0.4:
            alpha_lash = min(1.0, (ratio - 0.4) / 0.5)
            arc_y = cy + 1
            pts = np.array([
                [cx - rw, arc_y - 1],
                [cx - rw // 2, arc_y],
                [cx, arc_y + 1],
                [cx + rw // 2, arc_y],
                [cx + rw, arc_y - 1]
            ], np.int32).reshape((-1, 1, 2))
            lash_layer = res.copy()
            cv2.polylines(lash_layer, [pts], False, (40, 32, 35), 1, cv2.LINE_AA)
            res = cv2.addWeighted(lash_layer, alpha_lash, res, 1.0 - alpha_lash, 0)
            
    return res

def overlay_alpha_image(dst_bgr, src_bgra, x, y, scale=1.0, opacity=1.0):
    """Chèn sprite BGRA vào frame BGR với vị trí, tỷ lệ và độ mờ mong muốn"""
    if src_bgra is None or opacity <= 0.001:
        return dst_bgr
        
    if scale != 1.0:
        new_w = max(1, int(src_bgra.shape[1] * scale))
        new_h = max(1, int(src_bgra.shape[0] * scale))
        src_bgra = cv2.resize(src_bgra, (new_w, new_h), interpolation=cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LINEAR)
        
    sh, sw, _ = src_bgra.shape
    dh, dw, _ = dst_bgr.shape
    
    x1, y1 = max(0, int(x)), max(0, int(y))
    x2, y2 = min(dw, int(x + sw)), min(dh, int(y + sh))
    
    if x1 >= x2 or y1 >= y2:
        return dst_bgr
        
    sx1 = x1 - int(x)
    sy1 = y1 - int(y)
    sx2 = sx1 + (x2 - x1)
    sy2 = sy1 + (y2 - y1)
    
    src_crop = src_bgra[sy1:sy2, sx1:sx2]
    alpha = (src_crop[:, :, 3].astype(np.float32) / 255.0) * opacity
    alpha_3d = np.repeat(alpha[:, :, np.newaxis], 3, axis=2)
    
    dst_crop = dst_bgr[y1:y2, x1:x2].astype(np.float32)
    src_rgb = src_crop[:, :, :3].astype(np.float32)
    
    blended = src_rgb * alpha_3d + dst_crop * (1.0 - alpha_3d)
    dst_bgr[y1:y2, x1:x2] = np.clip(blended, 0, 255).astype(np.uint8)
    return dst_bgr

class BlinkingController:
    """Bộ điều khiển chu kỳ chớp mắt tự nhiên 3.5 - 4.5 giây / lần"""
    def __init__(self, cycle_sec=3.8):
        self.cycle_sec = cycle_sec

    def get_blink_ratio(self, t_sec):
        # Tính vị trí trong chu kỳ
        t_mod = t_sec % self.cycle_sec
        # Chớp mắt trong 0.12s cuối của chu kỳ (3 frame ở 30fps: 0.55 -> 1.0 -> 0.50)
        blink_start = self.cycle_sec - 0.14
        if t_mod >= blink_start:
            dt = t_mod - blink_start
            if dt < 0.04:
                return 0.55
            elif dt < 0.08:
                return 1.0
            elif dt < 0.12:
                return 0.50
        return 0.0

class AudioLipSyncController:
    """Phân tích sóng âm để lấy RMS và visemes chính xác"""
    def __init__(self, audio_file_path, fps=30, total_sec=180.0):
        self.fps = fps
        self.total_frames = int(total_sec * fps)
        self.frame_rms = np.zeros(self.total_frames, dtype=np.float32)
        self.speech_thresh = 0.02
        self.max_rms = 0.25
        
        if not audio_file_path or not os.path.exists(audio_file_path):
            self._setup_procedural()
            return
            
        wav_path = audio_file_path + ".temp_rms.wav"
        try:
            if not os.path.exists(wav_path) or os.path.getsize(wav_path) < 1000:
                cmd = ["ffmpeg", "-y", "-i", audio_file_path, "-f", "wav", "-ar", "44100", "-ac", "1", wav_path]
                subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
                
            wf = wave.open(wav_path, "rb")
            sr = wf.getframerate()
            n_frames = wf.getnframes()
            audio_bytes = wf.readframes(n_frames)
            wf.close()
            samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0
            samples_per_frame = sr / float(fps)
            
            for f_idx in range(self.total_frames):
                start_s = int(f_idx * samples_per_frame)
                end_s = int((f_idx + 1) * samples_per_frame)
                if start_s < len(samples):
                    chunk = samples[start_s:min(end_s, len(samples))]
                    self.frame_rms[f_idx] = np.sqrt(np.mean(chunk**2)) if len(chunk) > 0 else 0.0
                    
            self.max_rms = max(0.1, float(self.frame_rms.max()))
            self.speech_thresh = max(0.015, float(self.frame_rms.mean()) * 0.35)
            if os.path.exists(wav_path):
                os.remove(wav_path)
        except Exception:
            self._setup_procedural()

    def _setup_procedural(self):
        # Fallback nhịp nói tự nhiên
        self.frame_rms = np.array([0.08 if (i % 30 < 22 and (i // 90) % 4 != 3) else 0.0 for i in range(self.total_frames)], dtype=np.float32)
        self.speech_thresh = 0.03
        self.max_rms = 0.15

    def get_viseme(self, frame_idx, scene_end=False):
        if frame_idx >= len(self.frame_rms):
            return "smile" if scene_end else "closed"
            
        rms = self.frame_rms[frame_idx]
        if rms < self.speech_thresh:
            return "smile" if scene_end else "closed"
            
        energy_ratio = (rms - self.speech_thresh) / max(0.001, (self.max_rms - self.speech_thresh))
        step = (frame_idx // 3) % 4
        if energy_ratio > 0.45:
            return "ah" if step in (0, 2) else "oh"
        else:
            if step == 0:
                return "ah"
            elif step == 1:
                return "closed"
            elif step == 2:
                return "oh"
            else:
                return "ah"

    def get_head_nod_offset(self, frame_idx):
        if frame_idx >= len(self.frame_rms):
            return 0
        rms = self.frame_rms[frame_idx]
        if rms > self.speech_thresh:
            ratio = (rms - self.speech_thresh) / max(0.001, (self.max_rms - self.speech_thresh))
            return int(round(1.8 * min(1.0, ratio)))
        return 0

class SevenNguyenActor:
    """
    Lớp quản lý diễn xuất hoạt hình toàn diện của Host Seven Nguyen:
    Bao gồm load sprite, chớp mắt, khẩu hình, nhịp thở, gật đầu
    """
    def __init__(self, workspace_root):
        self.workspace = workspace_root
        output_dir = os.path.join(workspace_root, "output")
        assets_dir = os.path.join(workspace_root, "assets", "Seven Nguyen")
        
        # Load Sprites
        self.sprites = {}
        for k in ["p1", "p2", "p3", "p4", "p6"]:
            sp_path = os.path.join(output_dir, f"final_sprite_{k}.png")
            if os.path.exists(sp_path):
                self.sprites[k] = cv2.imread(sp_path, cv2.IMREAD_UNCHANGED)
            else:
                self.sprites[k] = None
                
        # Expressions Sheet
        exp_path = os.path.join(assets_dir, "seven_nguyen_animation_facial_expressions.jpg")
        if os.path.exists(exp_path):
            self.expressions_sheet = cv2.imread(exp_path)
        else:
            self.expressions_sheet = None
            
        # Trích xuất khẩu hình với alpha feathering
        self.mouth_patches = {}
        if self.expressions_sheet is not None:
            self.mouth_patches = {
                "closed": self._extract_tight_mouth(185, 217, rx=24, ry=12, feather=4),
                "smile":  self._extract_tight_mouth(520, 215, rx=25, ry=12, feather=4),
                "ah":     self._extract_tight_mouth(185, 595, rx=24, ry=14, feather=4),
                "oh":     self._extract_tight_mouth(520, 595, rx=20, ry=14, feather=4),
            }
            
        self.blink_ctl = BlinkingController(cycle_sec=3.8)
        
        # Cấu hình toạ độ đầu (mắt, miệng) cho từng pose
        self.pose_anchors = {
            "p1": {
                "mouth": {"cx": 180, "cy": 72, "scale": 0.38},
                "eyes":  {"lx": 168, "ly": 48, "rx": 192, "ry": 48, "rw": 5, "rh": 3}
            },
            "p2": {  # pointing right
                "mouth": {"cx": 150, "cy": 88, "scale": 0.40},
                "eyes":  {"lx": 136, "ly": 58, "rx": 162, "ry": 58, "rw": 5, "rh": 3}
            },
            "p3": {  # pointing left (Chuẩn chỉ vào bản đồ / card bên trái)
                "mouth": {"cx": 236, "cy": 88, "scale": 0.40},
                "eyes":  {"lx": 224, "ly": 58, "rx": 250, "ry": 58, "rw": 5, "rh": 3}
            },
            "p4": {  # crossed arms
                "mouth": {"cx": 95, "cy": 83, "scale": 0.42},
                "eyes":  {"lx": 78, "ly": 56, "rx": 108, "ry": 56, "rw": 6, "rh": 3}
            },
            "p6": {  # leaning / thinking
                "mouth": {"cx": 114, "cy": 106, "scale": 0.40},
                "eyes":  {"lx": 98, "ly": 76, "rx": 126, "ry": 76, "rw": 5, "rh": 3}
            }
        }

    def _extract_tight_mouth(self, cx, cy, rx=24, ry=14, feather=4):
        x1 = cx - rx - feather
        x2 = cx + rx + feather
        y1 = cy - ry - feather
        y2 = cy + ry + feather
        crop = self.expressions_sheet[y1:y2, x1:x2].copy()
        h, w = crop.shape[:2]
        
        Y, X = np.ogrid[:h, :w]
        dx = (X - (rx + feather)) / float(rx)
        dy = (Y - (ry + feather)) / float(ry)
        dist = np.sqrt(dx**2 + dy**2)
        
        outer = 1.0 + feather / float(min(rx, ry))
        alpha = np.zeros((h, w), dtype=np.float32)
        alpha[dist <= 1.0] = 1.0
        feather_mask = (dist > 1.0) & (dist < outer)
        t = (dist[feather_mask] - 1.0) / (outer - 1.0)
        alpha[feather_mask] = 0.5 * (1.0 + np.cos(np.pi * t))
        
        bgra = cv2.cvtColor(crop, cv2.COLOR_BGR2BGRA)
        bgra[:, :, 3] = np.clip(alpha * 255, 0, 255).astype(np.uint8)
        return bgra

    def _apply_mouth_patch(self, sprite_bgr, viseme_key, cx, cy, scale=0.40):
        if not self.mouth_patches:
            return sprite_bgr
        patch = self.mouth_patches.get(viseme_key, self.mouth_patches.get("closed"))
        if patch is None:
            return sprite_bgr
            
        pw = max(4, int(patch.shape[1] * scale))
        ph = max(4, int(patch.shape[0] * scale))
        p_sc = cv2.resize(patch, (pw, ph), interpolation=cv2.INTER_AREA)
        
        patch_bgr = p_sc[:, :, :3]
        mask = (p_sc[:, :, 3] > 80).astype(np.uint8) * 255
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.erode(mask, kernel, iterations=1)
        
        bh, bw = sprite_bgr.shape[:2]
        cx_safe = min(bw - pw // 2 - 2, max(pw // 2 + 2, int(cx)))
        cy_safe = min(bh - ph // 2 - 2, max(ph // 2 + 2, int(cy)))
        
        try:
            cloned = cv2.seamlessClone(patch_bgr, sprite_bgr[:, :, :3], mask, (cx_safe, cy_safe), cv2.NORMAL_CLONE)
            if sprite_bgr.shape[2] == 4:
                return np.dstack([cloned, sprite_bgr[:, :, 3]])
            return cloned
        except Exception:
            return sprite_bgr

    def get_procedural_viseme(self, t_sec, is_closing=False):
        """Khẩu hình sinh tự nhiên khi không có file audio RMS trực tiếp"""
        if is_closing:
            return "smile"
        # Chu kỳ câu: 3s nói, 0.4s nghỉ
        t_cycle = t_sec % 3.4
        if t_cycle > 3.0:
            return "closed"
        step = int(t_sec * 9.0) % 4
        if step == 0:
            return "ah"
        elif step == 1:
            return "closed"
        elif step == 2:
            return "oh"
        else:
            return "ah"

    def get_animated_sprite(self, pose_key, t_sec=0.0, frame_idx=0, audio_ctl=None, is_closing=False):
        """
        Tạo sprite có đầy đủ:
        - Khẩu hình lipsync
        - Chớp mắt tự nhiên
        - Gật đầu micro-nodding
        """
        raw_sprite = self.sprites.get(pose_key)
        if raw_sprite is None:
            raw_sprite = self.sprites.get("p1")
            if raw_sprite is None:
                return None
                
        sprite = raw_sprite.copy()
        anchor = self.pose_anchors.get(pose_key, self.pose_anchors["p1"])
        
        # 1. Viseme & Head nod
        if audio_ctl is not None:
            viseme = audio_ctl.get_viseme(frame_idx, scene_end=is_closing)
            nod_offset = audio_ctl.get_head_nod_offset(frame_idx)
        else:
            viseme = self.get_procedural_viseme(t_sec, is_closing=is_closing)
            nod_offset = int(round(1.5 * abs(math.sin(t_sec * 6.0)))) if viseme in ("ah", "oh") else 0
            
        # 2. Blinking ratio
        blink_ratio = self.blink_ctl.get_blink_ratio(t_sec)
        
        # 3. Apply mouth
        m_cfg = anchor["mouth"]
        sprite = self._apply_mouth_patch(sprite, viseme, m_cfg["cx"], m_cfg["cy"] + nod_offset, scale=m_cfg["scale"])
        
        # 4. Apply eyes blink
        e_cfg = anchor["eyes"]
        sprite_bgr = draw_natural_eye_blink(
            sprite[:, :, :3],
            e_cfg["lx"], e_cfg["ly"] + nod_offset,
            e_cfg["rx"], e_cfg["ry"] + nod_offset,
            rw=e_cfg["rw"], rh=e_cfg["rh"],
            ratio=blink_ratio
        )
        sprite[:, :, :3] = sprite_bgr
        return sprite

    def draw_actor(self, frame, pose_key, x, y, base_scale=1.0, opacity=1.0, t_sec=0.0, frame_idx=0, audio_ctl=None, is_closing=False, flip_horizontal=False):
        """
        Vẽ Seven Nguyen trực tiếp lên frame với nhịp thở tự nhiên, scale pulse, và các diễn hoạt mặt
        """
        sprite = self.get_animated_sprite(pose_key, t_sec=t_sec, frame_idx=frame_idx, audio_ctl=audio_ctl, is_closing=is_closing)
        if sprite is None:
            return frame
            
        if flip_horizontal:
            sprite = cv2.flip(sprite, 1)
            
        # Nhịp thở tự nhiên
        breath_y = math.sin(t_sec * 2.5 * math.pi) * 2.5
        breath_scale = base_scale * (1.0 + 0.003 * math.sin(t_sec * 2.5 * math.pi))
        
        return overlay_alpha_image(frame, sprite, int(x), int(y + breath_y), scale=breath_scale, opacity=opacity)