#!/usr/bin/env python3
"""
AVE lipsync — Kling LipSync 集成

功能:
  Audio→Video: 输入视频 + 语音音频 → 输出口型对齐的视频
  Text→Video: 输入视频 + 文本 → 输出口型对齐的视频 (TTS自动)

API:
  fal.ai endpoint （Kling LipSync 托管在 fal.ai）
  成本: $0.014/5s ≈ ¥0.10/秒 ≈ ¥0.50/5秒镜头

用法:
  python lipsync.py --video clip.mp4 --audio voice.wav --output synced.mp4
  python lipsync.py --video clip.mp4 --text "你好世界" --output synced.mp4

集成:
  作为口播策略后处理钩子, 在 video_factory.py 中调用
"""
import sys
import os
import json
import hashlib
import time
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from lib.config import load_config
from lib.cost_tracker import get_tracker
from lib.logger import get_logger

logger = get_logger("lipsync")

# ── 缓存 ──
CACHE_DIR = Path(os.environ.get("AVE_CACHE_DIR",
    str(Path.home() / "workbuddy-agent-os/agent-local/tools/ave/cache/lipsync")))
CACHE_DIR.mkdir(parents=True, exist_ok=True)

# ── fal.ai 端点 ──
# Kling LipSync 通过 fal.ai 提供
FAL_BASE = "https://fal.run"

ENDPOINTS = {
    "audio_to_video": f"{FAL_BASE}/fal-ai/kling-video/lipsync/audio-to-video",
    "text_to_video":  f"{FAL_BASE}/fal-ai/kling-video/lipsync/text-to-video",
}

# ═══════════════════════════════════════════════════════════
# 核心 API
# ═══════════════════════════════════════════════════════════

def lipsync_audio_to_video(
    video_path: str,
    audio_path: str,
    output_path: str = "",
    force: bool = False,
    api_key: str = "",
) -> str:
    """
    Audio→Video 唇形同步

    参数:
      video_path: 输入视频路径 (2-10s, 720p/1080p, ≤100MB)
      audio_path: 输入音频路径 (.mp3/.wav/.m4a, ≤5MB)
      output_path: 输出路径
      force: 强制重新生成
      api_key: fal.ai API Key

    返回:
      本地输出视频路径
    """
    if not output_path:
        joint = hashlib.md5(
            (video_path + audio_path).encode()
        ).hexdigest()[:16]
        output_path = str(CACHE_DIR / f"lipsync_{joint}.mp4")

    # 缓存
    if os.path.exists(output_path) and not force:
        logger.info(f"缓存命中: {output_path}")
        return output_path

    cfg = load_config()
    api_key = api_key or cfg.get("fal", {}).get("api_key", "")
    if not api_key:
        raise ValueError("缺少 fal.ai API Key (config fal.api_key)")

    logger.info(f"LipSync Audio→Video:")
    logger.info(f"  video: {video_path}")
    logger.info(f"  audio: {audio_path}")

    # 上传文件或传URL (fal.ai 支持直接传base64或文件URL)
    # 方案: 使用 fal.ai 的 upload 端点
    video_url = _upload_to_fal(video_path, api_key)
    audio_url = _upload_to_fal(audio_path, api_key)

    if not video_url or not audio_url:
        raise RuntimeError("上传文件到 fal.ai 失败")

    logger.info(f"  video_url: {video_url[:60]}...")
    logger.info(f"  audio_url: {audio_url[:60]}...")

    # 提交任务
    result = _submit_and_poll(
        endpoint=ENDPOINTS["audio_to_video"],
        payload={
            "video_url": video_url,
            "audio_url": audio_url,
        },
        api_key=api_key,
        timeout_sec=120,
    )

    # 下载结果
    video_url_res = result.get("video", {}).get("url", "") or result.get("output", "")
    if not video_url_res:
        raise RuntimeError(f"fal.ai 返回无视频URL: {json.dumps(result, ensure_ascii=False)[:200]}")

    logger.info(f"  下载结果中...")
    _download(video_url_res, output_path)

    # 费用
    tracker = get_tracker()
    try:
        dur = _get_media_duration(output_path)
    except Exception:
        dur = 5  # 默认5秒
    tracker.log("KlingLipSync", duration=dur, note=f"audio-to-video ¥0.014×{dur}s")

    sz_mb = os.path.getsize(output_path) / 1024 / 1024
    logger.info(f"  ✅ LipSync完成: {output_path} ({sz_mb:.1f}MB, {dur:.1f}s)")
    return output_path


def lipsync_text_to_video(
    video_path: str,
    text: str,
    voice_id: str = "default",
    output_path: str = "",
    force: bool = False,
    api_key: str = "",
) -> str:
    """
    Text→Video 唇形同步 (自动 TTS + 口型对齐, ≤120字)

    参数:
      video_path: 输入视频路径
      text: 口播文本 (≤120字)
      voice_id: 音色ID
      output_path: 输出路径
      force: 强制重新生成

    返回:
      本地输出视频路径
    """
    if len(text) > 120:
        logger.warning(f"文本过长 ({len(text)}字), 截断至120字")
        text = text[:120]

    if not output_path:
        joint = hashlib.md5(
            (video_path + text).encode()
        ).hexdigest()[:16]
        output_path = str(CACHE_DIR / f"lipsync_txt_{joint}.mp4")

    if os.path.exists(output_path) and not force:
        logger.info(f"缓存命中: {output_path}")
        return output_path

    cfg = load_config()
    api_key = api_key or cfg.get("fal", {}).get("api_key", "")
    if not api_key:
        raise ValueError("缺少 fal.ai API Key")

    logger.info(f"LipSync Text→Video: text={text[:40]}...")

    video_url = _upload_to_fal(video_path, api_key)
    if not video_url:
        raise RuntimeError("上传视频到 fal.ai 失败")

    result = _submit_and_poll(
        endpoint=ENDPOINTS["text_to_video"],
        payload={
            "video_url": video_url,
            "text": text,
            "voice_id": voice_id,
        },
        api_key=api_key,
        timeout_sec=180,
    )

    video_url_res = result.get("video", {}).get("url", "") or result.get("output", "")
    if not video_url_res:
        raise RuntimeError(f"fal.ai 返回无视频URL")

    _download(video_url_res, output_path)

    sz_mb = os.path.getsize(output_path) / 1024 / 1024
    logger.info(f"  ✅ LipSync Text→Video完成: {output_path} ({sz_mb:.1f}MB)")
    return output_path


# ═══════════════════════════════════════════════════════════
# 内部工具
# ═══════════════════════════════════════════════════════════

def _upload_to_fal(file_path: str, api_key: str) -> str:
    """上传文件到 fal.ai, 返回可访问的 URL"""
    import httpx
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")

    with open(file_path, "rb") as f:
        resp = httpx.post(
            f"{FAL_BASE}/fal-ai/upload",
            headers={
                "Authorization": f"Key {api_key}",
            },
            files={"file": f},
            timeout=120,
        )
        resp.raise_for_status()
        data = resp.json()
        return data.get("url", "")


def _submit_and_poll(
    endpoint: str,
    payload: dict,
    api_key: str,
    timeout_sec: int = 120,
) -> dict:
    """提交 fal.ai 任务并轮询, 返回结果"""
    import httpx
    headers = {
        "Authorization": f"Key {api_key}",
        "Content-Type": "application/json",
    }

    # 提交
    resp = httpx.post(
        endpoint.replace("/run", ""),  # 提交到非 /run 端点
        headers=headers,
        json=payload,
        timeout=30,
    )
    # fal.ai 的 submit API
    submit_url = endpoint
    resp = httpx.post(submit_url, headers=headers, json=payload, timeout=30)
    resp.raise_for_status()
    data = resp.json()

    # 检查是否直接返回结果（fal.ai 同步模式）
    if data.get("status") == "completed" or data.get("video"):
        return data

    # 获取请求ID进行轮询
    request_id = data.get("request_id", "")
    if not request_id:
        # 有的接口直接返回结果
        if data.get("video", {}).get("url"):
            return data
        raise RuntimeError(f"fal.ai 无 request_id: {json.dumps(data, ensure_ascii=False)[:200]}")

    status_url = f"{endpoint}/requests/{request_id}"
    logger.info(f"  轮询: {request_id}")

    start = time.time()
    for i in range(60):
        # 指数退避
        if i < 5:
            time.sleep(3)
        elif i < 15:
            time.sleep(5)
        else:
            time.sleep(10)

        elapsed = time.time() - start
        if elapsed > timeout_sec:
            raise TimeoutError(f"fal.ai 超时 ({timeout_sec}s)")

        try:
            q = httpx.get(status_url, headers=headers, timeout=15)
            q.raise_for_status()
            status_data = q.json()
            status = status_data.get("status", "")
            if status == "COMPLETED":
                return status_data
            elif status in ("FAILED", "ERROR"):
                raise RuntimeError(f"fal.ai 失败: {json.dumps(status_data, ensure_ascii=False)[:200]}")
            logger.debug(f"    [{i}] status={status}")
        except httpx.HTTPError as e:
            logger.debug(f"    [{i}] 查询失败: {e}")
            time.sleep(5)

    raise TimeoutError("fal.ai 轮询超时")


def _download(url: str, path: str):
    """下载文件"""
    import httpx
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with httpx.stream("GET", url, timeout=600) as r:
        r.raise_for_status()
        with open(path, "wb") as f:
            for chunk in r.iter_bytes(8192):
                f.write(chunk)


def _get_media_duration(path: str) -> float:
    """获取媒体时长 (秒) —— 用 lib.ffmpeg（本机 ffprobe 实为 ffmpeg 符号链接）"""
    from lib.ffmpeg import probe_duration
    return probe_duration(path)


# ═══════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════

def cli(args: list[str] | None = None):
    import argparse
    parser = argparse.ArgumentParser(description="Kling LipSync 唇形同步")
    sub = parser.add_subparsers(dest="mode", required=True)

    # audio → video
    p_av = sub.add_parser("audio-to-video", help="音频驱动唇形同步")
    p_av.add_argument("--video", required=True, help="输入视频")
    p_av.add_argument("--audio", required=True, help="输入音频")
    p_av.add_argument("--output", default="", help="输出路径")
    p_av.add_argument("--force", action="store_true", help="强制重新生成")

    # text → video
    p_tv = sub.add_parser("text-to-video", help="文本驱动唇形同步")
    p_tv.add_argument("--video", required=True, help="输入视频")
    p_tv.add_argument("--text", required=True, help="口播文本 (≤120字)")
    p_tv.add_argument("--voice-id", default="default", help="音色ID")
    p_tv.add_argument("--output", default="", help="输出路径")
    p_tv.add_argument("--force", action="store_true", help="强制重新生成")

    parsed = parser.parse_args(args)

    if parsed.mode == "audio-to-video":
        out = lipsync_audio_to_video(
            video_path=parsed.video,
            audio_path=parsed.audio,
            output_path=parsed.output,
            force=parsed.force,
        )
        print(f"\n✅ LipSync 完成: {out}")

    elif parsed.mode == "text-to-video":
        out = lipsync_text_to_video(
            video_path=parsed.video,
            text=parsed.text,
            voice_id=parsed.voice_id,
            output_path=parsed.output,
            force=parsed.force,
        )
        print(f"\n✅ LipSync 完成: {out}")


if __name__ == "__main__":
    cli()

# ════════════════════════════════════════════════════════════
# ★ 百炼 VideoRetalk（推荐实现，2026-10-03 实测探明）
#   文档：https://help.aliyun.com/zh/model-studio/videoretalk-api
#   优势：用现有百炼 key（无需 fal.ai）· ¥0.08/秒 · 1800 秒免费额度
#   上传复用 AVE 已有的两步式（getPolicy → OSS 直传 → oss://）
# ════════════════════════════════════════════════════════════

DASHSCOPE_BASE = "https://dashscope.aliyuncs.com"
VR_SUBMIT_URL = f"{DASHSCOPE_BASE}/api/v1/services/aigc/image2video/video-synthesis"
VR_TASK_URL = f"{DASHSCOPE_BASE}/api/v1/tasks/{{tid}}"
VIDEORETALK_MODEL = "videoretalk"
VR_PRICE_PER_SEC = 0.08          # 元/秒（官方后付费单价）
VR_FREE_QUOTA_SEC = 1800         # 官方免费额度（秒）


def vr_estimate_cost(duration_sec: float) -> float:
    """对口型成本估算（元）—— 实际优先消耗 1800 秒免费额度"""
    return round(max(0.0, float(duration_sec)) * VR_PRICE_PER_SEC, 2)


def _vr_upload(path: str, api_key: str) -> str:
    """两步式上传 → oss:// URL（与 person_swap 同一套；oss:// 需配 Resolve header）"""
    import httpx
    p = Path(path)
    r = httpx.get(f"{DASHSCOPE_BASE}/api/v1/uploads?action=getPolicy&model={VIDEORETALK_MODEL}",
                  headers={"Authorization": f"Bearer {api_key}"}, timeout=30)
    if r.status_code != 200:
        raise RuntimeError(f"获取上传凭证失败 HTTP {r.status_code}: {r.text[:200]}")
    d = r.json().get("data") or {}
    key = f"{d.get('upload_dir', '')}/{p.name}"
    form = {
        "OSSAccessKeyId": d.get("oss_access_key_id", ""),
        "policy": d.get("policy", ""),
        "signature": d.get("signature", ""),
        "key": key,
        "x-oss-object-acl": d.get("x_oss_object_acl", "private"),
        "x-oss-forbid-overwrite": "true",
    }
    with open(p, "rb") as f:
        r2 = httpx.post(d["upload_host"], data=form,
                        files={"file": (p.name, f, "application/octet-stream")}, timeout=300,
                        trust_env=False)
    if r2.status_code not in (200, 201, 204):
        raise RuntimeError(f"OSS 直传失败 HTTP {r2.status_code}: {r2.text[:200]}")
    return f"oss://{key}"


def lipsync_videoretalk(video_path: str, audio_path: str, output_path: str,
                        api_key: str = "", ref_image_path: str = "",
                        video_extension: bool = False,
                        timeout: int = 1800, progress_cb=None) -> str:
    """人物视频 + 人声 → **口型匹配**的新视频（百炼 VideoRetalk）

    · 成本 ¥0.08/秒（有 1800 秒免费额度）
    · 视频：正面近景、2~120s、mp4/mov、15~60fps、边长 640~2048
    · 音频：干净人声、2~120s（TTS 产物天然满足）
    · ref_image_path：角色参考图（可选，提升人脸匹配；我们正好有）

    返回输出视频本地路径。
    """
    import httpx
    api_key = api_key or (load_config().get("aliyun", {}) or {}).get("api_key", "")
    if not api_key:
        raise ValueError("VideoRetalk 需要百炼 API Key")
    _cb = progress_cb or (lambda p, n: None)

    vdur = _get_media_duration(video_path)
    if vdur and not (2.0 < vdur < 120.0):
        raise ValueError(f"视频时长 {vdur:.1f}s 超出 VideoRetalk 要求（2~120 秒）")

    _cb(5, "上传视频")
    vurl = _vr_upload(video_path, api_key)
    _cb(15, "上传音频")
    aurl = _vr_upload(audio_path, api_key)
    rurl = ""
    if ref_image_path and Path(ref_image_path).exists():
        try:
            _cb(20, "上传人脸参考图")
            rurl = _vr_upload(ref_image_path, api_key)
        except Exception as e:
            logger.warning("参考图上传失败（继续）：%s", e)

    inp = {"video_url": vurl, "audio_url": aurl}
    if rurl:
        inp["ref_image_url"] = rurl
    payload = {
        "model": VIDEORETALK_MODEL,
        "input": inp,
        "parameters": {"video_extension": bool(video_extension)},
    }
    hdr = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "X-DashScope-Async": "enable",
        "X-DashScope-OssResourceResolve": "enable",   # ★ 用 oss:// 时必须
    }
    _cb(25, "提交对口型任务")
    r = httpx.post(VR_SUBMIT_URL, headers=hdr, json=payload, timeout=60, trust_env=False)
    if r.status_code != 200:
        raise RuntimeError(f"提交失败 HTTP {r.status_code}: {r.text[:300]}")
    tid = (r.json().get("output") or {}).get("task_id", "")
    if not tid:
        raise RuntimeError(f"提交返回无 task_id: {r.text[:300]}")
    logger.info("  VideoRetalk 任务已提交: %s（预估 ¥%.2f）", tid, vr_estimate_cost(vdur))

    # 轮询
    deadline = time.time() + timeout
    last = ""
    while time.time() < deadline:
        rr = httpx.get(VR_TASK_URL.format(tid=tid),
                       headers={"Authorization": f"Bearer {api_key}"}, timeout=30, trust_env=False)
        out = (rr.json() or {}).get("output") or {}
        st = out.get("task_status", "")
        if st == "SUCCEEDED":
            url = out.get("video_url") or ""
            if not url:
                results = out.get("results") or []
                url = (results[0].get("video_url") or results[0].get("url")) if results else ""
            if not url:
                raise RuntimeError(f"成功但无视频 URL: {str(out)[:200]}")
            _cb(95, "下载结果")
            _download(url, output_path)
            _cb(100, "完成")
            logger.info("  ✅ 对口型完成: %s", Path(output_path).name)
            return output_path
        if st in ("FAILED", "CANCELED", "UNKNOWN"):
            raise RuntimeError(f"任务 {st}: {(out.get('message') or out.get('code') or '')[:200]}")
        if st != last:
            last = st
            logger.info("  [%s] 对口型任务状态: %s", tid[:8], st)
        _cb(30, f"处理中({st})")
        time.sleep(10)
    raise RuntimeError(f"对口型超时（>{timeout}s）")


# ── CLI 扩展：--provider videoretalk ──
def cli_videoretalk(video: str, audio: str, output: str,
                    ref_image: str = "", video_extension: bool = False):
    def _p(p, n):
        print(f"  {p}% {n}")
    out = lipsync_videoretalk(video, audio, output, ref_image_path=ref_image,
                              video_extension=video_extension, progress_cb=_p)
    print(f"\n✅ 输出: {out}")