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 DemucsTasks(LongTasks):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.gpu_id = int(os.environ.get('DEMUCS_GPU_ID', '5'))
    async def process_task(self, payload, workid=None):
        import json
        if isinstance(payload, str): payload = json.loads(payload)
        task_type = payload.get('task_type', '')
        if task_type == 'separate':
            from workers.separate import run_separate
            return await run_separate(self, payload)
        raise ValueError(f'Unknown task_type: {task_type}')

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

def init():
    env = ServerEnv()
    env.longtasks = DemucsTasks('redis://127.0.0.1:6379', 'demucs', worker_cnt=1, stuck_seconds=600, max_age_hours=24)
    add_startup(on_app_built)

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