#!/usr/bin/env python3
"""
Sophnet Video Generation Script
Supports text-to-video and image-to-video; default model is Seedance unless the user specifies another model.
"""
import os
import sys
import subprocess
import argparse
import json
import time
import requests
from typing import Optional
import sophnet_tools

# API Configuration
SOPH_API_KEY = sophnet_tools.get_api_key()
if SOPH_API_KEY is None:
    print("错误: 未找到 Sophnet API Key，请确保已正确配置环境变量 SOPH_API_KEY")
    sys.exit(1)
GENERATE_URL = "https://www.sophnet.com/api/open-apis/projects/easyllms/videogenerator/generate"

def create_request(
    prompt: Optional[str],
    model: str = "Seedance-2.0",
    negative_prompt: Optional[str] = None,
    first_frame_url: Optional[str] = None,
    last_frame_url: Optional[str] = None,
    size: str = "1280x720",
    duration: int = 5,
    generate_audio: bool = True,
    draft: bool = False
) -> dict:
    """Build request body for video generation."""
    content = []

    # Add text content
    if prompt:
        text_content = {"type": "text", "text": prompt}
        if negative_prompt:
            text_content["negative_prompt"] = negative_prompt
        content.append(text_content)

    # Add first frame
    if first_frame_url:
        content.append({
            "type": "image_url",
            "image_url": {"url": first_frame_url},
            "role": "first_frame"
        })
        # Modality fix only: same family T2V + image → I2V (does not replace unrelated user model choices).
        if model == "happyhorse-1.0-t2v":
            model = "happyhorse-1.0-i2v"
        elif model == "Wan2.6-T2V":
            model = "Wan2.6-I2V"

        # Add last frame
        if last_frame_url:
            content.append({
                "type": "image_url",
                "image_url": {"url": last_frame_url},
                "role": "last_frame"
            })

    # Build parameters.
    # Seedance uses resolution levels (e.g. 720p/1080p), while some other models accept size (e.g. 1280*720).
    def infer_resolution_from_size(raw_size: str) -> str:
        normalized = raw_size.lower().replace("*", "x")
        if normalized in {"1920x1080", "1080x1920"}:
            return "1080p"
        if normalized in {"1280x720", "720x1280", "1088x832", "832x1088", "960x960"}:
            return "720p"
        return "720p"

    parameters = {"duration": duration}
    if model in {"Seedance-2.0", "Seedance-2.0-Fast", "Seedance-1.5-Pro"}:
        parameters["resolution"] = infer_resolution_from_size(size)
    else:
        parameters["size"] = size.replace("x", "*")

    return {"model": model, "content": content, "parameters": parameters}


def make_request(req_data: dict) -> dict:
    """Make HTTP request with authentication."""
    headers = {
        "Authorization": f"Bearer {SOPH_API_KEY}",
        "Content-Type": "application/json"
    }
    try:
        response = requests.request(
            "POST",
            GENERATE_URL,
            data=json.dumps(req_data).encode("utf-8"),
            headers=headers,
            timeout=60
        )
        if response.status_code == 200:
            return response.json()
        else:
            print(f"HTTP请求失败，状态码: {response.status_code}, 响应内容: {response.text}")
            return None
    except requests.exceptions.RequestException as e:
        print(f"HTTP请求超时: {e}")
        return None


def make_get_request(task_id: str) -> dict:
    """Make GET request with authentication."""
    headers = {
        "Authorization": f"Bearer {SOPH_API_KEY}",
        "Content-Type": "application/json"
    }
    url = f"{GENERATE_URL}/{task_id}"
    response = requests.request("GET", url, headers=headers, timeout=60)
    if response.status_code == 200:
        return response.json().get("result", {})
    else:
        print(f"HTTP请求失败，状态码: {response.status_code}, 响应内容: {response.text}")
        return None


def is_url(path: str) -> bool:
    """Check if the input string is a URL."""
    return path.startswith(("http://", "https://"))


# 使用浏览器式请求头，避免 CDN/OSS 因 User-Agent 拒绝下载
DOWNLOAD_HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
    "Accept": "*/*",
    "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
}


def get_file_with_requests(url: str, result_name: str, result_path: str) -> bool:
    """使用 requests 下载文件（带浏览器请求头，避免被服务器拒绝）。"""
    try:
        out_dir = os.path.expanduser(result_path)
        os.makedirs(out_dir, exist_ok=True)
        out_path = os.path.join(out_dir, result_name)
        resp = requests.get(url, headers=DOWNLOAD_HEADERS, timeout=30, stream=True)
        resp.raise_for_status()
        with open(out_path, "wb") as f:
            for chunk in resp.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        return os.path.isfile(out_path) and os.path.getsize(out_path) > 0
    except requests.exceptions.RequestException as e:
        print(f"requests 下载失败: {e}")
        return False
    except Exception as e:
        print(f"下载过程中发生异常: {e}")
        return False


def get_file_with_wget(url: str, result_name: str, result_path: str) -> bool:
    """Download file using wget (with browser User-Agent)."""
    try:
        os.makedirs(os.path.expanduser(result_path), exist_ok=True)
        out_file = os.path.join(os.path.expanduser(result_path), result_name)
        user_agent = DOWNLOAD_HEADERS["User-Agent"]
        result = subprocess.run(
            ["wget", "-O", out_file,
             "--tries=1", "--timeout=30",
             f"--user-agent={user_agent}",
             url],
            capture_output=True,
            text=True
        )
        if result.returncode == 0 and os.path.isfile(out_file) and os.path.getsize(out_file) > 0:
            return True
        else:
            print(f"wget 下载失败，返回码: {result.returncode}, 标准输出: {result.stdout}, 标准错误: {result.stderr}")
            return False
    except Exception as e:
        print(f"执行 wget 过程中发生异常: {e}")
        return False


def gen_video(
    prompt: str,
    model: str = "Seedance-2.0",
    size: str = "1280*720",
    duration: int = 5,
    first_frame: Optional[str] = None,
) -> str:
    """Generate video using Sophnet Video Generation API."""
    result_video_url = None
    input_image = None
    local_video_url = None
    result_path = "/home/node/.openclaw/workspace/media/inbound/videos/"
    # Handle first frame
    if first_frame:
        if is_url(first_frame):
            input_image = first_frame
        elif os.path.isfile(first_frame):
            input_image = sophnet_tools.upload_oss(first_frame) 
            if input_image is None:
                print("上传第一帧图片失败")
                sys.exit(1)
    # Build request
    request_body = create_request(
        model=model,
        prompt=prompt,
        first_frame_url=input_image,
        size=size,
        duration=duration,
    )
    print(json.dumps(request_body, indent=2))

    request = make_request(request_body)
    if request is None:
        return result_video_url, local_video_url  
    if request.get("status") != 0:
        print(f"视频生成失败，错误信息: {request.get('message', '未知错误')}")
        return result_video_url, local_video_url
    task_id = request.get("result", {}).get("task_id")
    if not task_id:
        print("视频生成失败，未返回任务ID")
        return result_video_url
    print(f"视频生成请求已提交，任务ID: {task_id}")  
    # Poll for completion
    time_start = time.time()
    while True:
        request = make_get_request(task_id)
        if request is None:
            print("获取任务状态失败")
            return result_video_url, local_video_url
        
        if request.get("status") == "succeeded":
            video_url = request.get("content", {}).get("video_url")
            if video_url:
                if get_file_with_requests(video_url, result_name=f"{task_id}.mp4", result_path=result_path):
                    local_video_url = os.path.join(os.path.expanduser(result_path), f"{task_id}.mp4")
                    # Re-upload to get a persistent URL
                    result_video_url = sophnet_tools.upload_oss(local_video_url)
                    if result_video_url:
                        return result_video_url, local_video_url
                    else:
                        return video_url, local_video_url
                else:
                    return video_url, local_video_url
            else:
                print(f"视频生成失败")
                return None, local_video_url  # Return the direct URL if download fails
        else:
            print(f"当前状态: {request}")
            if request.get("status", "").find("failed") != -1:
                print(f"视频生成失败，错误信息: {request.get('status', '未知错误')}")
                return result_video_url, local_video_url          
            print(f"视频生成中，当前状态: {request.get('status', '未知状态')}")           
            if time.time() - time_start > 300:  # 5分钟超时
                print("视频生成超时")
                return result_video_url, local_video_url       
            time.sleep(5)

if __name__ == "__main__":
    # Listed for UX hints only; unknown models are still sent to the API as the user passed them.
    supported_models = [
        "happyhorse-1.0-t2v",
        "happyhorse-1.0-i2v",
        "Wan2.6-T2V",
        "Wan2.6-I2V",
        "ViduQ2-pro",
        "Seedance-2.0",
        "Seedance-2.0-Fast",
        "Seedance-1.5-Pro",
    ]
    parser = argparse.ArgumentParser(
        description="Generate videos using Sophnet Video Generation API"
    )
    parser.add_argument("--prompt", help="Text prompt (required for text-to-video)")
    parser.add_argument("--model", type=str, default="Seedance-2.0",
                       help="Model name. Defaults to Seedance-2.0")
    parser.add_argument("--size", type=str, default="1280*720", help="Resolution (default: 1280*720)")
    parser.add_argument("--duration", type=int, default=5, help="Video duration in seconds (default: 5)")
    parser.add_argument("--first-frame", type=str, help="First frame image URL or path")
    parser.add_argument("--generate-audio", action="store_true", help="Generate audio")
    args = parser.parse_args()

    model = args.model
    # HappyHorse default: T2V + image → I2V (API expects I2V when a frame is supplied).
    if args.first_frame and model == "happyhorse-1.0-t2v":
        model = "happyhorse-1.0-i2v"
    # Friendly name normalization only (canonical Vidu / Seedance names).
    elif model.find("ViduQ2") != -1:
        model = "ViduQ2-pro"
    elif model.strip().lower() == "seedance":
        model = "Seedance-2.0"
    elif model.strip().lower() == "seedance-fast":
        model = "Seedance-2.0-Fast"
    elif model not in supported_models:
        print(f"提示: 模型 '{args.model}' 不在内置常用列表中，将按您指定的名称原样调用 API。")
    
    result_url, local_video_url = gen_video(
        args.prompt,
        model,
        args.size,
        args.duration,
        args.first_frame,
    )
    
    if result_url:
        print(f"\n您的视频提示词为: {args.prompt}")
        print(f"视频生成成功: {result_url}")
        print(f"视频本地保存路径: {local_video_url}")
        print(f"提示: 视频保存24小时，请及时下载")
        
        # 生成支持的模型列表（排除当前使用的模型）
        other_models = [m for m in supported_models if m != model]
        other_models_str = "、".join(other_models)
        print(f"当前使用的模型为 {model}，还支持 {other_models_str}")
        print(f"如果想要指定模型，可以直接对我说：使用 [模型名称] 模型，生成 [视频描述]")
        print(
            "提示: 常用模型示例: happyhorse-1.0-t2v、happyhorse-1.0-i2v、"
            "Wan2.6-T2V、Wan2.6-I2V、ViduQ2-pro、Seedance-2.0、Seedance-2.0-Fast、Seedance-1.5-Pro（也可传入文档中的其他模型名）"
        )
    else:
        print("\n视频生成失败")
        sys.exit(1)
