youtube-auto-dub/web_app.py

#!/usr/bin/env python3
"""Guardio web UI for launching YouTube Auto Dub jobs."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
import os
import subprocess
import sys
import threading
import uuid

from flask import Flask, abort, jsonify, render_template, request, send_from_directory

from main import build_parser
from src.audio_separation import DEFAULT_MIX_MODE
from src.engines import OUTPUT_DIR


BASE_DIR = Path(__file__).resolve().parent
LOG_DIR = BASE_DIR / "logs" / "guardio"


@dataclass
class DubJob:
    """Runtime state for a web-launched dub job."""

    id: str
    command: list[str]
    log_path: Path
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    status: str = "queued"
    returncode: int | None = None
    completed_at: datetime | None = None


JOBS: dict[str, DubJob] = {}
JOBS_LOCK = threading.Lock()


def _utc_iso(value: datetime | None) -> str | None:
    if value is None:
        return None
    return value.astimezone(timezone.utc).isoformat()


def build_pipeline_command(form: dict[str, str]) -> list[str]:
    """Build a validated command for the existing CLI pipeline."""
    parser = build_parser()
    args = parser.parse_args(_form_to_cli_args(form))
    command = [
        sys.executable,
        str(BASE_DIR / "main.py"),
        args.url,
        "--lang",
        args.lang,
        "--mix-mode",
        args.mix_mode,
        "--translation-backend",
        args.translation_backend,
    ]

    optional_flags = {
        "--browser": args.browser,
        "--cookies": args.cookies,
        "--whisper_model": args.whisper_model,
        "--lmstudio-base-url": args.lmstudio_base_url,
        "--lmstudio-model": args.lmstudio_model,
    }
    for flag, value in optional_flags.items():
        if value:
            command.extend([flag, value])

    if args.gpu:
        command.append("--gpu")

    return command


def _form_to_cli_args(form: dict[str, str]) -> list[str]:
    url = (form.get("url") or "").strip()
    if not url:
        raise ValueError("A YouTube URL is required.")

    cli_args = [url]
    field_flags = {
        "lang": "--lang",
        "browser": "--browser",
        "cookies": "--cookies",
        "whisper_model": "--whisper_model",
        "mix_mode": "--mix-mode",
        "translation_backend": "--translation-backend",
        "lmstudio_base_url": "--lmstudio-base-url",
        "lmstudio_model": "--lmstudio-model",
    }

    defaults = {
        "lang": "es",
        "mix_mode": DEFAULT_MIX_MODE,
        "translation_backend": "lmstudio",
    }

    for field_name, flag in field_flags.items():
        value = (form.get(field_name) or defaults.get(field_name) or "").strip()
        if value:
            cli_args.extend([flag, value])

    if form.get("gpu") in {"1", "true", "on", "yes"}:
        cli_args.append("--gpu")

    return cli_args


def _serialize_job(job: DubJob) -> dict[str, object]:
    return {
        "id": job.id,
        "status": job.status,
        "returncode": job.returncode,
        "created_at": _utc_iso(job.created_at),
        "completed_at": _utc_iso(job.completed_at),
        "log": _read_log_tail(job.log_path),
    }


def _read_log_tail(log_path: Path, max_chars: int = 20000) -> str:
    if not log_path.exists():
        return ""
    text = log_path.read_text(encoding="utf-8", errors="replace")
    return text[-max_chars:]


def _run_job(job: DubJob) -> None:
    with JOBS_LOCK:
        job.status = "running"

    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"

    with job.log_path.open("w", encoding="utf-8", errors="replace") as log_file:
        log_file.write("Guardio started a YouTube Auto Dub job.\n")
        log_file.write(f"Command: {' '.join(job.command)}\n\n")
        log_file.flush()

        process = subprocess.Popen(
            job.command,
            cwd=BASE_DIR,
            env=env,
            stdout=log_file,
            stderr=subprocess.STDOUT,
            text=True,
        )
        returncode = process.wait()

    with JOBS_LOCK:
        job.returncode = returncode
        job.completed_at = datetime.now(timezone.utc)
        job.status = "succeeded" if returncode == 0 else "failed"


def _list_outputs() -> list[dict[str, object]]:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    files = sorted(
        (path for path in OUTPUT_DIR.glob("*") if path.is_file()),
        key=lambda path: path.stat().st_mtime,
        reverse=True,
    )
    return [
        {
            "name": path.name,
            "size_mb": round(path.stat().st_size / (1024 * 1024), 1),
            "modified_at": datetime.fromtimestamp(path.stat().st_mtime, timezone.utc).isoformat(),
        }
        for path in files[:20]
    ]


def create_app() -> Flask:
    """Create the Guardio Flask app."""
    app = Flask(__name__)

    @app.get("/")
    def index():
        return render_template("index.html", outputs=_list_outputs())

    @app.get("/api/jobs")
    def list_jobs():
        with JOBS_LOCK:
            jobs = sorted(JOBS.values(), key=lambda item: item.created_at, reverse=True)
            return jsonify([_serialize_job(job) for job in jobs])

    @app.post("/api/jobs")
    def create_job():
        try:
            command = build_pipeline_command(request.form)
        except (SystemExit, ValueError) as exc:
            return jsonify({"error": str(exc) or "Invalid job options."}), 400

        LOG_DIR.mkdir(parents=True, exist_ok=True)
        job_id = uuid.uuid4().hex[:12]
        job = DubJob(
            id=job_id,
            command=command,
            log_path=LOG_DIR / f"{job_id}.log",
        )

        with JOBS_LOCK:
            JOBS[job.id] = job

        thread = threading.Thread(target=_run_job, args=(job,), daemon=True)
        thread.start()
        return jsonify(_serialize_job(job)), 202

    @app.get("/api/jobs/<job_id>")
    def get_job(job_id: str):
        with JOBS_LOCK:
            job = JOBS.get(job_id)
            if job is None:
                abort(404)
            return jsonify(_serialize_job(job))

    @app.get("/api/outputs")
    def list_outputs():
        return jsonify(_list_outputs())

    @app.get("/outputs/<path:filename>")
    def download_output(filename: str):
        return send_from_directory(OUTPUT_DIR, filename, as_attachment=True)

    return app


app = create_app()


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=7860, debug=True)