from __future__ import annotations

import asyncio
import os
import sys
import uuid
import wave
from collections import deque
from pathlib import Path

import cv2
import numpy as np

from config import settings


class AvatarService:
    def __init__(self, width: int = 640, height: int = 360) -> None:
        self.width = width
        self.height = height
        self._frames: deque[np.ndarray] = deque()
        self._lock = asyncio.Lock()
        self._job_lock = asyncio.Lock()
        self._last_error: str | None = None
        self._base_frame = self._load_base_frame()

    async def render_frame(self, mouth_open: float = 0.0, speaking: bool = False) -> np.ndarray:
        frame = self._base_frame.copy()
        if speaking:
            alpha = float(np.clip(mouth_open, 0.0, 1.0))
            # Slightly boost contrast/brightness while speaking to keep a "live" feeling.
            frame = cv2.convertScaleAbs(frame, alpha=1.0 + 0.10 * alpha, beta=2 + int(8 * alpha))
            # Fallback mouth animation on realistic base frame.
            h, w = frame.shape[:2]
            cx = w // 2
            cy = int(h * 0.70)
            half_w = max(28, int(w * 0.055))
            half_h = max(4, int(4 + alpha * h * 0.035))

            lip_color = (26, 26, 92)
            lip_border = (70, 70, 160)
            cv2.ellipse(frame, (cx, cy), (half_w + 2, half_h + 2), 0, 0, 360, lip_border, -1)
            cv2.ellipse(frame, (cx, cy), (half_w, half_h), 0, 0, 360, lip_color, -1)
        return frame

    async def idle_frame(self) -> np.ndarray:
        return await self.render_frame(mouth_open=0.0, speaking=False)

    async def speaking_frame(self, mouth_open: float = 0.5) -> np.ndarray:
        return await self.render_frame(mouth_open=mouth_open, speaking=True)

    async def pop_generated_frame(self) -> np.ndarray | None:
        async with self._lock:
            if self._frames:
                return self._frames.popleft()
        return None

    async def enqueue_musetalk_from_audio(self, audio: np.ndarray, sr: int) -> None:
        if not settings.musetalk_enabled:
            return
        if audio.size <= 1 or sr <= 0:
            return

        async with self._job_lock:
            try:
                job_id = f"job_{uuid.uuid4().hex[:10]}"
                repo_dir = Path(settings.musetalk_repo_dir)
                work_root = repo_dir / "results" / "visual-chat-jobs" / job_id
                work_root.mkdir(parents=True, exist_ok=True)
                wav_path = work_root / "input.wav"
                cfg_path = work_root / "inference.yaml"
                output_name = f"{job_id}.mp4"
                output_path = Path(settings.musetalk_result_dir) / "v15" / output_name

                self._write_wav(wav_path, audio.astype(np.float32, copy=False), sr)
                cfg_path.write_text(
                    (
                        "task_0:\n"
                        f"  video_path: \"{settings.musetalk_source_video}\"\n"
                        f"  audio_path: \"{wav_path}\"\n"
                        f"  result_name: \"{output_name}\"\n"
                    ),
                    encoding="utf-8",
                )

                env = os.environ.copy()
                env.setdefault("HF_HOME", str(repo_dir / "models"))
                env["PYTHONPATH"] = (
                    f"{repo_dir}:{env['PYTHONPATH']}" if env.get("PYTHONPATH") else str(repo_dir)
                )
                proc = await asyncio.create_subprocess_exec(
                    sys.executable,
                    str(settings.musetalk_infer_script),
                    "--version",
                    "v15",
                    "--inference_config",
                    str(cfg_path),
                    "--result_dir",
                    str(settings.musetalk_result_dir),
                    "--unet_model_path",
                    str(settings.musetalk_unet_model),
                    "--unet_config",
                    str(settings.musetalk_unet_config),
                    "--whisper_dir",
                    str(settings.musetalk_whisper_dir),
                    "--batch_size",
                    "8",
                    "--use_float16",
                    cwd=str(repo_dir),
                    env=env,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                _, stderr = await proc.communicate()
                if proc.returncode != 0:
                    self._last_error = stderr.decode("utf-8", errors="ignore")[-500:]
                    return
                frames: list[np.ndarray] = []
                if output_path.exists():
                    frames = self._read_video_frames(str(output_path))
                if not frames:
                    # Some MuseTalk runs leave image frames instead of final mp4.
                    fallback_dir = Path(settings.musetalk_result_dir) / "v15" / "yongen"
                    if fallback_dir.exists():
                        frames = self._read_image_frames(str(fallback_dir))
                if not frames:
                    self._last_error = f"MuseTalk produced no readable frames: {output_path}"
                    return

                async with self._lock:
                    self._frames.clear()
                    self._frames.extend(frames)
                self._last_error = None
            except Exception as exc:  # pragma: no cover
                self._last_error = str(exc)

    def _read_video_frames(self, path: str) -> list[np.ndarray]:
        cap = cv2.VideoCapture(path)
        frames: list[np.ndarray] = []
        while True:
            ok, frame = cap.read()
            if not ok or frame is None:
                break
            if frame.shape[1] != self.width or frame.shape[0] != self.height:
                frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_LINEAR)
            frames.append(frame)
        cap.release()
        return frames

    def _read_image_frames(self, dir_path: str) -> list[np.ndarray]:
        frames: list[np.ndarray] = []
        for name in sorted(os.listdir(dir_path)):
            if not name.lower().endswith(".png"):
                continue
            frame = cv2.imread(os.path.join(dir_path, name))
            if frame is None:
                continue
            if frame.shape[1] != self.width or frame.shape[0] != self.height:
                frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_LINEAR)
            frames.append(frame)
        return frames

    def _load_base_frame(self) -> np.ndarray:
        src = settings.musetalk_source_video
        if src and os.path.exists(src):
            cap = cv2.VideoCapture(src)
            ok, frame = cap.read()
            cap.release()
            if ok and frame is not None:
                if frame.shape[1] != self.width or frame.shape[0] != self.height:
                    frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_LINEAR)
                return frame
        # Last-resort dark background (should rarely happen).
        fallback = np.zeros((self.height, self.width, 3), dtype=np.uint8)
        fallback[:, :, :] = (20, 20, 20)
        return fallback

    @staticmethod
    def _write_wav(path: Path, audio: np.ndarray, sr: int) -> None:
        pcm = np.clip(audio * 32767.0, -32768, 32767).astype(np.int16)
        with wave.open(str(path), "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(sr)
            wf.writeframes(pcm.tobytes())

    @property
    def health(self) -> dict:
        return {
            "enabled": settings.musetalk_enabled,
            "queued_frames": len(self._frames),
            "last_error": self._last_error,
        }
