# -*- coding:utf-8 -*-
"""
Wan2.2-TI2V-5B 视频生成 worker
调用官方 generate.py 脚本
"""
import os
import json
import uuid
import asyncio
import subprocess
from datetime import datetime
from appPublic.log import debug, exception

OUTPUT_DIR = '/data/ymq/wan22-outputs'
REPO_DIR = '/data/ymq/wan22-service/repo'
MODEL_PATH = '/data/ymq/models/Wan-AI/Wan2.2-TI2V-5B'
PYTHON = '/share/vllm-0.8.5/bin/python'


async def run_generate(longtasks, payload):
    """
    Execute video generation via generate.py subprocess.
    
    payload: {
        task_id: str,
        prompt: str,
        image: str (optional path for i2v),
        size: str (default "1280*720"),
        frame_num: int (default 81, must be 4n+1),
        sample_steps: int (optional),
        sample_guide_scale: float (optional),
        base_seed: int (optional),
    }
    """
    task_id = payload.get('task_id', str(uuid.uuid4())[:12])
    prompt = payload.get('prompt', '')
    image = payload.get('image', None)
    size = payload.get('size', '1280*720')
    frame_num = payload.get('frame_num', 81)
    sample_steps = payload.get('sample_steps', None)
    sample_guide_scale = payload.get('sample_guide_scale', None)
    base_seed = payload.get('base_seed', None)
    
    # Ensure frame_num is 4n+1
    frame_num = max(17, min(frame_num, 129))
    if (frame_num - 1) % 4 != 0:
        frame_num = ((frame_num - 1) // 4) * 4 + 1
    
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    output_file = os.path.join(OUTPUT_DIR, f'{task_id}.mp4')
    
    # Build command
    cmd = [
        PYTHON, 'generate.py',
        '--task', 'ti2v-5B',
        '--ckpt_dir', MODEL_PATH,
        '--size', size,
        '--frame_num', str(frame_num),
        '--prompt', prompt,
        '--save_file', output_file,
        '--offload_model', 'True',
    ]
    
    if image:
        cmd.extend(['--image', image])
    
    if sample_steps:
        cmd.extend(['--sample_steps', str(sample_steps)])
    
    if sample_guide_scale:
        cmd.extend(['--sample_guide_scale', str(sample_guide_scale)])
    
    if base_seed is not None:
        cmd.extend(['--base_seed', str(base_seed)])
    
    # Set CUDA_VISIBLE_DEVICES for single GPU
    gpu_id = longtasks.gpu_id if longtasks.gpu_id else 2
    env = os.environ.copy()
    env['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
    
    debug(f'Running: {" ".join(cmd)}')
    
    try:
        # Run in subprocess
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=REPO_DIR,
            env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        stdout, stderr = await proc.communicate()
        
        if proc.returncode != 0:
            error_msg = stderr.decode('utf-8', errors='ignore')[-500:]
            exception(f'generate.py failed: {error_msg}')
            return {
                'task_id': task_id,
                'status': 'failed',
                'error': error_msg
            }
        
        # Check output file
        if not os.path.exists(output_file):
            return {
                'task_id': task_id,
                'status': 'failed',
                'error': 'Output file not created'
            }
        
        file_size = os.path.getsize(output_file)
        
        return {
            'task_id': task_id,
            'status': 'completed',
            'video_url': f'/idfile?path={task_id}.mp4',
            'video_path': output_file,
            'size': size,
            'frame_num': frame_num,
            'file_size': file_size,
            'prompt': prompt[:100]
        }
        
    except Exception as e:
        exception(f'Generation error: {e}')
        return {
            'task_id': task_id,
            'status': 'failed',
            'error': str(e)
        }