# -*- coding:utf-8 -*-
import os
from ahserver.webapp import webapp
from ahserver.serverenv import ServerEnv
from ahserver.configuredServer import add_startup
from longtasks.longtasks import LongTasks, schedule_once
from appPublic.log import debug


class Wan22Tasks(LongTasks):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.gpu_id = int(os.environ.get('WAN22_GPU_ID', '2'))

    async def process_task(self, payload: dict, workid: int = None):
        import json
        if isinstance(payload, str):
            payload = json.loads(payload)
        task_type = payload.get('task_type', '')
        debug(f'Wan22Tasks processing: type={task_type}')
        if task_type == 'generate_video':
            from workers.generate import run_generate
            return await run_generate(self, payload)
        raise ValueError(f'Unknown task_type: {task_type}')


async def on_app_built(app):
    env = ServerEnv()
    longtasks = env.longtasks
    if longtasks:
        schedule_once(0.1, longtasks.run)
        debug(f'longtasks worker started, GPU: {longtasks.gpu_id}')


def init():
    env = ServerEnv()
    longtasks = Wan22Tasks(
        'redis://127.0.0.1:6379',
        'wan22',
        worker_cnt=1,
        stuck_seconds=3600,
        max_age_hours=24
    )
    env.longtasks = longtasks
    add_startup(on_app_built)


if __name__ == '__main__':
    webapp(init)