#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
=============================================================================
THẬP NGŨ NIÊN (THE FIFTEEN SPRINGS) - MUSE2API CLIENT ENGINE
=============================================================================
Module tích hợp chính thức cho cổng dịch vụ Muse2API Gateway (FastAPI):
- Kết nối trực tiếp Base URL: http://127.0.0.1:18610 (OpenAI API Compatible /v1)
- Tự động nạp API Key từ c:\\Projects\\Muse2API\\data\\api_key hoặc biến môi trường
- Render video I2V / T2V bất đồng bộ (Async Task Pattern: POST /v1/videos -> Poll)
- Tự động mã hóa Start Frame sang Base64 Data URI (lossless PNG/JPEG)
- Áp dụng Audio Guard chuẩn mực và kiểm duyệt từ khóa nhạy cảm
- Tự động xác định versioning (_v1, _v2...), tải file MP4, xử lý reverse motion,
  tích hợp âm thanh chuyên biệt và trích xuất Tail Frame cho Head-Tail Chaining.
=============================================================================
"""

import os
import sys
import re
import time
import base64
import mimetypes
import argparse
import subprocess
import shutil
from pathlib import Path
from typing import Optional, Tuple, Dict, Any

import cv2
import requests

# Đảm bảo UTF-8 Console Windows
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

BASE_DIR = Path(__file__).resolve().parent.parent
ASSETS_DIR = BASE_DIR / "04_Assets"
VIDEOS_DIR = ASSETS_DIR / "videos"
KEYFRAMES_DIR = ASSETS_DIR / "keyframes"

LOCAL_KEY_FILE = Path(r"c:\Projects\Muse2API\data\api_key")
DEFAULT_FALLBACK_KEY = "m2a--nqA-RT1mZ4AWalrdDVzSYkhI9Cksd8O"
DEFAULT_BASE_URL = "http://127.0.0.1:18610"

AUDIO_GUARD_CANONICAL = " Quy tắc âm thanh: Tuyệt đối KHÔNG sinh nhạc nền (no music/BGM), không âm thanh điện tử, không tạp âm rè nhiễu. Chỉ sinh âm thanh môi trường tự nhiên (foley, ambience) và thoại nhân vật chân thực."


def get_ffmpeg() -> str:
    """Xác định đường dẫn nhị phân FFmpeg."""
    exe = Path(sys.executable).parent / "ffmpeg.exe"
    if exe.exists():
        return str(exe)
    if shutil.which("ffmpeg"):
        return "ffmpeg"
    return "ffmpeg"


def sanitize_prompt(prompt: str) -> str:
    """Thanh lọc từ khóa nhạy cảm kích hoạt bộ lọc kiểm duyệt Meta/Muse."""
    if not prompt:
        return prompt
    replacements = [
        (r'\b(nấm mồ|mồ mả|nghĩa địa|mộ địa|mộ phần|ngôi mộ|ngôi mả)\b', 'gò đất phủ cỏ úa xưa cũ'),
        (r'\b(âm khí|tử khí|âm khí nặng nề)\b', 'gió lạnh chiều tà se sắt hoang vu'),
        (r'\b(khóc người đời xưa|khóc than bạc mệnh)\b', 'nghiêng mình trầm tư đượm buồn'),
        (r'\b(rải vàng mã|tiền vàng mã|vàng mã)\b', 'thắp nén hương trầm phảng phất'),
        (r'\b(xương tàn|xương cốt|xác chết|chết chóc)\b', 'dấu tích xưa cũ'),
        (r'\b(bạc mệnh)\b', 'số phận long đong'),
    ]
    cleaned = prompt
    for pattern, repl in replacements:
        cleaned = re.sub(pattern, repl, cleaned, flags=re.IGNORECASE)
    return cleaned


def ensure_audio_guard(prompt: str) -> str:
    """Đảm bảo prompt tuân thủ chỉ thị Audio Guard cấm nhạc nền."""
    if not prompt:
        return AUDIO_GUARD_CANONICAL.strip()
    prompt = sanitize_prompt(prompt)
    if "Tuyệt đối KHÔNG sinh nhạc nền" not in prompt and "Strictly NO background music" not in prompt:
        return prompt.rstrip() + AUDIO_GUARD_CANONICAL
    return prompt


def resolve_api_key() -> str:
    """Đọc API Key theo thứ tự ưu tiên: Env Var -> Local Key File -> Fallback Default."""
    env_key = os.environ.get("MUSE_API_KEY")
    if env_key and env_key.strip():
        return env_key.strip()

    if LOCAL_KEY_FILE.exists():
        try:
            content = LOCAL_KEY_FILE.read_text(encoding="utf-8").strip()
            if content:
                return content
        except Exception:
            pass

    return DEFAULT_FALLBACK_KEY


def encode_image_to_data_uri(image_path: str) -> str:
    """Chuyển đổi tệp ảnh cục bộ thành Base64 Data URI."""
    p = Path(image_path).resolve()
    if not p.is_file():
        raise FileNotFoundError(f"Không tìm thấy file ảnh: {image_path}")

    mime, _ = mimetypes.guess_type(str(p))
    if not mime:
        ext = p.suffix.lower()
        if ext in (".jpg", ".jpeg"):
            mime = "image/jpeg"
        elif ext == ".png":
            mime = "image/png"
        elif ext == ".webp":
            mime = "image/webp"
        else:
            mime = "application/octet-stream"

    raw_bytes = p.read_bytes()
    b64_str = base64.b64encode(raw_bytes).decode("ascii")
    return f"data:{mime};base64,{b64_str}"


class MuseApiClient:
    """Client giao tiếp trực tiếp với cổng Muse2API Gateway."""

    def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None):
        self.base_url = (base_url or os.environ.get("MUSE_API_BASE_URL") or DEFAULT_BASE_URL).rstrip("/")
        self.api_key = api_key or resolve_api_key()
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def check_health(self, timeout: float = 3.0) -> Dict[str, Any]:
        """Kiểm tra sức khỏe cổng dịch vụ và trạng thái bể tài khoản (Account Pool)."""
        url = f"{self.base_url}/readyz"
        resp = requests.get(url, timeout=timeout)
        resp.raise_for_status()
        return resp.json()

    def is_available(self, timeout: float = 2.0) -> bool:
        """Kiểm tra nhanh xem Muse2API có đang phản hồi không."""
        try:
            data = self.check_health(timeout=timeout)
            return bool(data.get("ready", False))
        except Exception:
            return False

    def create_video_task(
        self,
        prompt: str,
        image: Optional[str] = None,
        model: str = "sora",
        size: str = "16:9",
        duration: int = 5,
        timeout: float = 30.0
    ) -> str:
        """
        Gửi yêu cầu tạo video (Text-to-Video hoặc Image-to-Video).
        Trả về task_id (vd: vid_xxxx).
        """
        payload = {
            "model": model,
            "prompt": ensure_audio_guard(prompt),
            "size": size,
            "duration": duration
        }

        if image:
            if image.startswith(("data:", "http://", "https://")):
                payload["image"] = image
            else:
                # File cục bộ -> Encode Data URI
                payload["image"] = encode_image_to_data_uri(image)

        url = f"{self.base_url}/v1/videos"
        resp = requests.post(url, headers=self.headers, json=payload, timeout=timeout)
        resp.raise_for_status()
        data = resp.json()
        task_id = data.get("id")
        if not task_id:
            raise RuntimeError(f"Muse2API không trả về task_id hợp lệ: {data}")
        return task_id

    def poll_video_task(
        self,
        task_id: str,
        timeout: float = 1200.0,
        poll_interval: float = 4.0,
        verbose: bool = True
    ) -> Dict[str, Any]:
        """
        Lặp theo dõi tiến độ tạo video cho đến khi hoàn thành hoặc lỗi.
        Trả về thông tin kết quả chứa URL video.
        """
        url = f"{self.base_url}/v1/videos/{task_id}"
        started = time.time()
        last_progress = -1

        while time.time() - started < timeout:
            time.sleep(poll_interval)
            try:
                resp = requests.get(url, headers=self.headers, timeout=15.0)
                resp.raise_for_status()
                data = resp.json()
            except Exception as e:
                if verbose:
                    print(f"      [~] Tạm thời mất kết nối tới gateway ({e}), đang thử lại...")
                continue

            status = data.get("status")
            progress = data.get("progress", 0)

            if verbose and progress != last_progress:
                print(f"      [API Task {task_id}] Trạng thái: {status.upper()} ({progress}%)")
                last_progress = progress

            if status == "succeeded":
                result = data.get("result") or {}
                video_url = result.get("url")
                if not video_url:
                    raise RuntimeError(f"Task thành công nhưng thiếu URL video: {data}")
                return data

            if status == "failed":
                err = data.get("error")
                raise RuntimeError(f"Tạo video thất bại từ Muse2API: {err}")

        raise TimeoutError(f"Hết thời gian chờ tạo video ({timeout}s) cho task {task_id}.")

    def download_media(self, media_url: str, output_path: Path, timeout: float = 120.0) -> Path:
        """Tải video MP4 từ Muse2API và lưu vào ổ cứng."""
        output_path.parent.mkdir(parents=True, exist_ok=True)
        resp = requests.get(media_url, headers=self.headers, stream=True, timeout=timeout)
        resp.raise_for_status()

        temp_file = output_path.with_name(f"{output_path.name}.downloading")
        with open(temp_file, "wb") as f:
            for chunk in resp.iter_content(chunk_size=65536):
                if chunk:
                    f.write(chunk)

        if temp_file.exists() and temp_file.stat().st_size > 1000:
            shutil.move(str(temp_file), str(output_path))
            return output_path
        else:
            if temp_file.exists():
                temp_file.unlink()
            raise RuntimeError(f"File tải về từ {media_url} bị rỗng hoặc không hợp lệ.")


def get_next_video_path(shot_id: str) -> Path:
    """Tự động xác định tên version tiếp theo (_v1, _v2, _v3...)."""
    VIDEOS_DIR.mkdir(parents=True, exist_ok=True)
    pattern = re.compile(rf"^{re.escape(shot_id)}_10s(?:_v(\d+))?\.mp4$", re.IGNORECASE)
    highest_v = 0
    has_unversioned = False
    for f in VIDEOS_DIR.glob(f"{shot_id}_10s*.mp4"):
        m = pattern.match(f.name)
        if m:
            v_num = m.group(1)
            if v_num is not None:
                highest_v = max(highest_v, int(v_num))
            else:
                has_unversioned = True
    next_v = max(highest_v + 1, 2 if has_unversioned and highest_v == 0 else highest_v + 1)
    return VIDEOS_DIR / (f"{shot_id}_10s_v{next_v}.mp4" if (highest_v > 0 or has_unversioned) else f"{shot_id}_10s.mp4")


def render_shot_via_api(
    shot_id: str,
    prompt: str,
    input_frame: Optional[str] = None,
    reverse_motion: bool = False,
    duration: int = 5,
    size: str = "16:9",
    client: Optional[MuseApiClient] = None
) -> Tuple[bool, Optional[str]]:
    """
    Quy trình render hoàn chỉnh cho 1 Shot thông qua Muse2API Gateway:
    1. Kiểm tra kết nối gateway
    2. Gửi tác vụ POST /v1/videos (kèm First-Frame Base64 nếu có)
    3. Giám sát polling đến khi hoàn tất
    4. Tải video MP4 về 04_Assets/videos/<shot_id>_10s_v<N>.mp4
    5. Đảo ngược chuyển động FFmpeg (nếu reverse_motion=True)
    6. Muxing âm thanh chuyên biệt (nếu có SFX / Voice)
    7. Trích xuất Tail Frame (clean_frame_239.jpg) phục vụ Head-Tail Chaining.
    """
    if client is None:
        client = MuseApiClient()

    print(f"      ⚡ [Muse2API] Bắt đầu gọi API tới: {client.base_url} (Key: {client.api_key[:6]}***)...")

    # Kiểm tra health
    try:
        health = client.check_health(timeout=3.0)
        acc_info = health.get("accounts", {})
        print(f"      ✓ Gateway trực tuyến: {acc_info.get('available', 0)}/{acc_info.get('total', 0)} tài khoản khả dụng.")
    except Exception as e:
        print(f"      [!] Lỗi kết nối tới Muse2API Gateway: {e}")
        return False, None

    # Xác định đường dẫn file đầu ra
    target_video_path = get_next_video_path(shot_id)
    print(f"      🎯 Đích xuất bản video: {target_video_path.name}")

    if input_frame:
        p_frame = Path(input_frame)
        if p_frame.exists():
            print(f"      🖼️ Start Frame: {p_frame.name} ({p_frame.stat().st_size // 1024} KB)")
        else:
            print(f"      ⚠️ Start Frame không tồn tại: {input_frame}, chuyển sang Text-to-Video.")
            input_frame = None

    # 1. Gửi task
    try:
        task_id = client.create_video_task(
            prompt=prompt,
            image=input_frame,
            model="sora",
            size=size,
            duration=duration
        )
        print(f"      ✓ Đã khởi tạo tác vụ: {task_id}, bắt đầu theo dõi...")
    except Exception as e:
        print(f"      [!] Không thể khởi tạo task video: {e}")
        return False, None

    # 2. Polling
    try:
        task_data = client.poll_video_task(task_id, timeout=1200.0, poll_interval=4.0)
        video_url = task_data["result"]["url"]
        print(f"      ✓ Render hoàn tất! Đang tải từ {video_url}...")
    except Exception as e:
        print(f"      [!] Quá trình render video thất bại: {e}")
        return False, None

    # 3. Download
    try:
        client.download_media(video_url, target_video_path)
        print(f"      ✓ Đã tải và lưu video vào: {target_video_path}")
    except Exception as e:
        print(f"      [!] Không thể tải video kết quả: {e}")
        return False, None

    # 4. Reverse Motion nếu được yêu cầu
    if reverse_motion:
        print(f"      🔄 Đang đảo ngược chuyển động video qua FFmpeg...")
        temp_reversed = target_video_path.with_name(f"{target_video_path.stem}_rev_temp.mp4")
        ffmpeg_bin = get_ffmpeg()
        cmd_rev = [
            ffmpeg_bin, "-y",
            "-i", str(target_video_path),
            "-vf", "reverse",
            "-af", "areverse",
            "-c:v", "libx264", "-crf", "18", "-preset", "slow",
            "-c:a", "aac", "-b:a", "192k", "-ar", "48000",
            str(temp_reversed)
        ]
        res_rev = subprocess.run(cmd_rev, capture_output=True, text=True)
        if res_rev.returncode == 0 and temp_reversed.exists() and temp_reversed.stat().st_size > 1000:
            shutil.move(str(temp_reversed), str(target_video_path))
            print(f"      ✓ Đã đảo ngược video và audio thành công!")
        elif temp_reversed.exists():
            temp_reversed.unlink()

    # 5. Tích hợp âm thanh SFX / Voice nếu có
    audio_cand = None
    for a_dir in [ASSETS_DIR / "audio_sfx", ASSETS_DIR / "audio_voice", ASSETS_DIR / "audio"]:
        if not a_dir.exists():
            continue
        for ext in [".wav", ".m4a", ".mp3", ".aac"]:
            af = a_dir / f"{shot_id}{ext}"
            if af.exists():
                audio_cand = af
                break
        if audio_cand:
            break

    if audio_cand:
        print(f"      🎵 Phát hiện âm thanh chuyên biệt: {audio_cand.name}. Đang tích hợp vào video...")
        temp_muxed = target_video_path.with_name(f"{target_video_path.stem}_mux_temp.mp4")
        ffmpeg_bin = get_ffmpeg()
        cmd_mux = [
            ffmpeg_bin, "-y",
            "-i", str(target_video_path),
            "-i", str(audio_cand),
            "-c:v", "copy",
            "-c:a", "aac", "-b:a", "192k", "-ar", "48000",
            "-shortest",
            str(temp_muxed)
        ]
        res_mux = subprocess.run(cmd_mux, capture_output=True, text=True)
        if res_mux.returncode == 0 and temp_muxed.exists() and temp_muxed.stat().st_size > 1000:
            shutil.move(str(temp_muxed), str(target_video_path))
            print(f"      ✓ Đã gắn âm thanh vào video thành công!")
        elif temp_muxed.exists():
            temp_muxed.unlink()

    # 6. Trích xuất Tail Frame (clean_frame_239.jpg)
    print(f"      📸 Đang trích xuất Tail Frame cho Head-Tail Chaining...")
    shot_kf_dir = KEYFRAMES_DIR / shot_id
    shot_kf_dir.mkdir(parents=True, exist_ok=True)
    tail_frame_path = shot_kf_dir / "clean_frame_239.jpg"

    cap = cv2.VideoCapture(str(target_video_path))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, total_frames - 1))
    ret, tail_frame = cap.read()
    cap.release()

    if ret and tail_frame is not None:
        cv2.imwrite(str(tail_frame_path), tail_frame)
        print(f"      ✓ Đã trích xuất Tail Frame: {tail_frame_path}")
    else:
        print(f"      ⚠️ Không thể trích xuất tail frame bằng OpenCV, thử dùng FFmpeg...")
        ffmpeg_bin = get_ffmpeg()
        cmd_tf = [
            ffmpeg_bin, "-y",
            "-sseof", "-0.1",
            "-i", str(target_video_path),
            "-vframes", "1",
            "-q:v", "2",
            str(tail_frame_path)
        ]
        subprocess.run(cmd_tf, capture_output=True)

    print(f"\n[✓] HOÀN TẤT THÀNH CÔNG (Muse2API): {target_video_path.name} -> Tail: {tail_frame_path.name}\n")
    return True, str(tail_frame_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Muse2API Client - Kiểm tra và render video qua API")
    parser.add_argument("--health", action="store_true", help="Kiểm tra trạng thái gateway và pool tài khoản")
    parser.add_argument("--shot", dest="shot_id", default=None, help="Mã định danh shot cần render (vd: ep01_scene02_shot04)")
    parser.add_argument("--prompt", default=None, help="Prompt mô tả chuyển động")
    parser.add_argument("--input-frame", default=None, help="Đường dẫn Start Frame")
    parser.add_argument("--reverse", action="store_true", help="Kích hoạt reverse motion")
    parser.add_argument("--duration", type=int, default=5, help="Thời lượng video (giây)")
    args = parser.parse_args()

    client = MuseApiClient()

    if args.health:
        print(f"Kiểm tra kết nối tới Muse2API: {client.base_url}...")
        try:
            h = client.check_health()
            import json
            print(json.dumps(h, indent=2, ensure_ascii=False))
        except Exception as e:
            print(f"Lỗi: {e}")
            sys.exit(1)
        sys.exit(0)

    if args.shot_id:
        p = args.prompt
        f = args.input_frame
        if not p or not f:
            try:
                from production_orchestrator import get_all_shots, resolve_start_frame
                all_s = get_all_shots()
                s_data = all_s.get(args.shot_id, {})
                if not p:
                    p = s_data.get("motion_prompt")
                if not f:
                    f = resolve_start_frame(args.shot_id, s_data)
            except Exception:
                pass

        if not p:
            print(f"Lỗi: Không tìm thấy prompt cho shot: {args.shot_id}")
            sys.exit(1)

        ok, tail = render_shot_via_api(
            shot_id=args.shot_id,
            prompt=p,
            input_frame=f,
            reverse_motion=args.reverse,
            duration=args.duration,
            client=client
        )
        if not ok:
            sys.exit(1)