SonicForgeStudio/app/tasks/worker.py

import os
import uuid
import time
import glob
import logging
from app.config import settings
from app.core.analyzer import analyze_audio, analyze_structure_with_ai
from app.core.audio_editor import (
    edit_audio_file, cut_and_loop_segment, mix_multitrack_session, export_audio
)
from app.core.dsp_utils import find_nearest_zero_crossing_file

logger = logging.getLogger(__name__)

# ──────────────────────────────────────────────────────────────────────────
# Task layer 2 che do:
#  - Server/Docker: celery day du (broker Redis) — dung nhu cu.
#  - Desktop slim (PyInstaller KHONG bundle celery/redis): task chay in-process
#    (thread nen + registry dict), API contract GIONG het (.delay() tra
#    task_id, /tasks/{id} tra status/result) nen frontend khong doi gi.
# ──────────────────────────────────────────────────────────────────────────
try:
    from celery import Celery
    HAS_CELERY = True
except Exception:  # pragma: no cover - frozen desktop slim build
    Celery = None
    HAS_CELERY = False

if HAS_CELERY:
    celery_app = Celery(
        "audio_tasks",
        broker=settings.CELERY_BROKER_URL,
        backend=settings.CELERY_RESULT_BACKEND
    )

    celery_app.conf.update(
        task_serializer="json",
        accept_content=["json"],
        result_serializer="json",
        timezone="UTC",
        enable_utc=True,
    )

    # Che do desktop (SF_DESKTOP=1, do desktop_engine.py set): chay task dong bo
    # trong tien trinh (eager) — ban Standalone KHONG kem Redis broker.
    if os.getenv("SF_DESKTOP") == "1":
        celery_app.conf.update(
            task_always_eager=True,
            task_eager_propagates=True,
            broker_url="memory://",
            result_backend="cache+memory://",
        )

    # ── Lich trinh tu dong don dep file het han (Week 5) ──
    celery_app.conf.beat_schedule = {
        "cleanup-expired-files-every-hour": {
            "task": "app.tasks.worker.cleanup_expired_files_task",
            "schedule": 3600.0,  # Chay moi gio
        },
    }
else:
    celery_app = None
    # Registry in-process cho desktop slim: task_id -> {"status", "result"/"error"}
    _results = {}


def _task(fn):
    """Wrapper: celery task (server) hoac in-process task (desktop slim)."""
    if HAS_CELERY:
        return celery_app.task(fn)
    return _InProcessTask(fn)


class _InProcessTask:
    """Task chay tren thread nen, ket qua luu vao registry dict — dung cho
    bundle desktop khong kem celery (tiet kiem ~40MB)."""

    def __init__(self, fn):
        self._fn = fn

    def delay(self, *args, **kwargs):
        import threading
        tid = uuid.uuid4().hex
        _results[tid] = {"status": "PENDING"}

        def _run():
            try:
                result = self._fn(*args, **kwargs)
                _results[tid] = {"status": "SUCCESS", "result": result}
            except Exception as e:  # noqa: BLE001 - bao loi day du cho UI
                logger.exception("In-process task %s failed", tid)
                _results[tid] = {"status": "FAILURE", "error": str(e)}

        threading.Thread(target=_run, daemon=True, name=f"task-{tid[:8]}").start()
        return _SimpleAsyncResult(tid)


class _SimpleAsyncResult:
    """Giong celery.result.AsyncResult ve mat API cho desktop slim."""

    def __init__(self, task_id):
        self.task_id = task_id

    @property
    def id(self):
        """Giong celery.result.AsyncResult.id — audio.py dung task.id."""
        return self.task_id

    @property
    def status(self):
        return _results.get(self.task_id, {}).get("status", "PENDING")

    @property
    def result(self):
        return _results.get(self.task_id, {}).get("result")

    def ready(self):
        return _results.get(self.task_id, {}).get("status") in ("SUCCESS", "FAILURE")

    def successful(self):
        return self.status == "SUCCESS"


def get_task_result(task_id):
    """Tra AsyncResult (celery) hoac _SimpleAsyncResult (desktop slim)."""
    if HAS_CELERY:
        from celery.result import AsyncResult
        return AsyncResult(task_id, app=celery_app)
    return _SimpleAsyncResult(task_id)


@_task
def analyze_audio_task(file_id: str):
    file_path = os.path.join(settings.UPLOADS_DIR, file_id)
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Upload file not found: {file_id}")
    return analyze_audio(file_path)


@_task
def analyze_ai_task(file_id: str, api_base_url: str = None,
                    model: str = "deepseek-chat"):
    """
    Task phân tích cấu trúc khuôn nhạc bằng AI (Week 4).
    API key is read from OPENAI_API_KEY env var only (never serialized into task queue).
    """
    upload_path = os.path.join(settings.UPLOADS_DIR, file_id)
    processed_path = os.path.join(settings.PROCESSED_DIR, file_id)
    
    if os.path.exists(processed_path):
        file_path = processed_path
    elif os.path.exists(upload_path):
        file_path = upload_path
    else:
        raise FileNotFoundError(f"File not found: {file_id}")
    
    return analyze_structure_with_ai(
        file_path=file_path,
        api_base_url=api_base_url,
        model=model
    )


@_task
def edit_audio_task(config: dict):
    file_id = config.get("file_id")
    
    input_path = os.path.join(settings.UPLOADS_DIR, file_id)
    if not os.path.exists(input_path):
        input_path = os.path.join(settings.PROCESSED_DIR, file_id)
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Source file not found: {file_id}")
            
    output_path = os.path.join(settings.PROCESSED_DIR, file_id)
    
    res = edit_audio_file(config, input_path, output_path)
    return {
        "file_id": file_id,
        "success": True,
        "details": res
    }


@_task
def export_audio_task(file_id: str, format: str = "wav",
                      sample_rate: int = 44100, bit_depth: int = 16):
    """
    Task xuất tệp âm thanh sang nhiều định dạng (WAV/MP3/OGG).
    """
    upload_path = os.path.join(settings.UPLOADS_DIR, file_id)
    processed_path = os.path.join(settings.PROCESSED_DIR, file_id)
    
    if os.path.exists(processed_path):
        source_path = processed_path
    elif os.path.exists(upload_path):
        source_path = upload_path
    else:
        raise FileNotFoundError(f"File not found: {file_id}")
    
    # Tạo output filename
    base_name = os.path.splitext(file_id)[0]
    output_filename = f"{base_name}_exported.{format}"
    output_path = os.path.join(settings.PROCESSED_DIR, output_filename)
    
    result = export_audio(
        input_path=source_path,
        output_path=output_path,
        format=format,
        sample_rate=sample_rate,
        bit_depth=bit_depth
    )
    
    result["output_file_id"] = output_filename
    return result


@_task
def mix_multitrack_task(session_config: dict):
    """
    Task xử lý hòa âm đa kênh (Multitrack Mixdown).
    """
    session_id = session_config.get("session_id")
    export_settings = session_config.get("export_settings", {})
    tracks = session_config.get("tracks", [])
    
    # Chuẩn bị metadata cho từng track
    tracks_meta = []
    for track in tracks:
        if track.get("muted", False):
            continue
            
        file_id = track.get("file_id")
        
        # Tìm file nguồn
        upload_path = os.path.join(settings.UPLOADS_DIR, file_id)
        processed_path = os.path.join(settings.PROCESSED_DIR, file_id)
        
        if os.path.exists(processed_path):
            file_path = processed_path
        elif os.path.exists(upload_path):
            file_path = upload_path
        else:
            raise FileNotFoundError(f"File not found: {file_id}")
            
        tracks_meta.append({
            "file_path": file_path,
            "volume": track.get("volume", 1.0),
            "muted": False
        })
    
    # Tạo tên file output
    output_format = export_settings.get("format", "wav")
    output_filename = f"{session_id}_mixed.{output_format}"
    output_path = os.path.join(settings.PROCESSED_DIR, output_filename)
    
    # Thực hiện mix
    result = mix_multitrack_session(
        tracks_meta=tracks_meta,
        output_path=output_path,
        sample_rate=export_settings.get("sample_rate", 44100),
        bit_depth=export_settings.get("bit_depth", 16)
    )
    
    result["output_file_id"] = output_filename
    result["session_id"] = session_id
    
    return result


@_task
def process_multitrack_session_task(session_config: dict):
    """
    Task xử lý toàn bộ session với nhiều tracks và clips.
    Xử lý từng clip (với zero-crossing alignment), sau đó hòa âm tất cả tracks.
    """
    session_id = session_config.get("session_id")
    export_settings = session_config.get("export_settings", {})
    tracks = session_config.get("tracks", [])
    
    processed_tracks = []
    
    # Xử lý từng track
    for track in tracks:
        if track.get("muted", False):
            continue
            
        track_id = track.get("track_id")
        file_id = track.get("file_id")
        clips = track.get("clips", [])
        
        # Tìm file nguồn
        upload_path = os.path.join(settings.UPLOADS_DIR, file_id)
        processed_path = os.path.join(settings.PROCESSED_DIR, file_id)
        
        if os.path.exists(processed_path):
            source_path = processed_path
        elif os.path.exists(upload_path):
            source_path = upload_path
        else:
            raise FileNotFoundError(f"File not found: {file_id}")
        
        # Xử lý clips nếu có
        if clips:
            clip = clips[0]
            
            start_sec = clip.get("start_time_seconds")
            end_sec = clip.get("end_time_seconds")
            
            # Áp dụng zero-crossing alignment nếu được yêu cầu
            if clip.get("apply_zero_crossing", True):
                start_sec = find_nearest_zero_crossing_file(source_path, start_sec)
                end_sec = find_nearest_zero_crossing_file(source_path, end_sec)
            
            # Sử dụng hàm cut_and_loop_segment
            processed_segment = cut_and_loop_segment(
                file_path=source_path,
                start_sec=start_sec,
                end_sec=end_sec,
                loop_count=clip.get("loop_count", 1),
                fade_in_ms=clip.get("fade_in_ms", 150),
                fade_out_ms=clip.get("fade_out_ms", 150),
                volume_db_change=0.0
            )
            
            # Lưu segment đã xử lý
            temp_filename = f"temp_{track_id}_{uuid.uuid4().hex[:8]}.wav"
            temp_path = os.path.join(settings.PROCESSED_DIR, temp_filename)
            processed_segment.export(temp_path, format="wav")
            
            track_file_path = temp_path
        else:
            track_file_path = source_path
        
        processed_tracks.append({
            "file_path": track_file_path,
            "volume": track.get("volume", 1.0),
            "muted": False
        })
    
    # Hòa âm tất cả tracks
    output_format = export_settings.get("format", "wav")
    output_filename = f"{session_id}_final.{output_format}"
    output_path = os.path.join(settings.PROCESSED_DIR, output_filename)
    
    result = mix_multitrack_session(
        tracks_meta=processed_tracks,
        output_path=output_path,
        sample_rate=export_settings.get("sample_rate", 44100),
        bit_depth=export_settings.get("bit_depth", 16)
    )
    
    # Dọn dẹp các file tạm
    for track in processed_tracks:
        if "temp_" in os.path.basename(track["file_path"]):
            try:
                os.remove(track["file_path"])
            except Exception as e:
                logger.warning("Failed to remove temp file %s: %s", track["file_path"], e)
    
    result["output_file_id"] = output_filename
    result["session_id"] = session_id
    
    return result


@_task
def cleanup_expired_files_task(max_age_hours: int = 24):
    """
    Task tự động dọn dẹp các tệp kết xuất hết hạn (Week 5).
    Xóa file trong thư mục processed cũ hơn max_age_hours giờ.
    """
    now = time.time()
    max_age_seconds = max_age_hours * 3600
    
    cleaned_count = 0
    cleaned_size = 0
    
    for directory in [settings.PROCESSED_DIR]:
        if not os.path.exists(directory):
            continue
            
        for filepath in glob.glob(os.path.join(directory, "*")):
            if os.path.isfile(filepath):
                file_age = now - os.path.getmtime(filepath)
                if file_age > max_age_seconds:
                    file_size = os.path.getsize(filepath)
                    try:
                        os.remove(filepath)
                        cleaned_count += 1
                        cleaned_size += file_size
                    except Exception as e:
                        logger.warning("Failed to remove expired file %s: %s", filepath, e)
    
    return {
        "cleaned_files": cleaned_count,
        "cleaned_size_mb": round(cleaned_size / (1024 * 1024), 2),
        "max_age_hours": max_age_hours
    }


@_task
def render_project_task(project_id: str, project_name: str, project_json_str: str, sample_rate: int = 44100):
    """
    Task Celery để kết xuất dự án ngoại tuyến (Offline Project Mixdown) áp dụng specs 30_DAW_ARCHITECT.md.
    """
    import json
    from app.core.render_engine import PythonRenderEngine
    
    project_json = json.loads(project_json_str)
    output_filename = f"{project_id}_render.wav"
    output_path = os.path.join(settings.PROCESSED_DIR, output_filename)
    
    engine = PythonRenderEngine(sample_rate=sample_rate)
    engine.render_project(project_json, output_path)
    
    return {
        "project_id": project_id,
        "project_name": project_name,
        "success": True,
        "output_file_id": output_filename
    }