"""Stem worker: polls MySQL for queued jobs, runs Demucs with the chosen options, saves stems as tracks.
pip install demucs pymysql   (GPU strongly recommended)"""
import os, time, glob, json, shutil, subprocess, pymysql
STORE = os.environ.get("STORE", "../backend/storage")
db = pymysql.connect(host=os.environ.get("DB_HOST","127.0.0.1"), user=os.environ.get("DB_USER","root"),
    password=os.environ.get("DB_PASS",""), database="studio", autocommit=True, cursorclass=pymysql.cursors.DictCursor)
while True:
    with db.cursor() as c:
        c.execute("SELECT * FROM stem_jobs WHERE status='queued' ORDER BY id LIMIT 1")
        job = c.fetchone()
        if not job: time.sleep(3); continue
        c.execute("UPDATE stem_jobs SET status='running' WHERE id=%s", (job["id"],))
        c.execute("SELECT * FROM tracks WHERE id=%s", (job["source_track_id"],))
        src = c.fetchone()
    out = f"/tmp/demucs_{job['id']}"
    try:
        o = json.loads(job.get("options") or "{}")
        model, two, keep = o.get("model", "htdemucs"), o.get("two"), o.get("keep") or []
        cmd = ["demucs", "-n", model, "-o", out] + (["--two-stems", two] if two else []) + [os.path.join(STORE, src["file_path"])]
        subprocess.run(cmd, check=True)
        res = []
        for f in glob.glob(f"{out}/*/*/*.wav"):
            stem = os.path.splitext(os.path.basename(f))[0]
            if keep and stem not in keep: continue
            fn = f"stem_{job['id']}_{stem}.wav"; shutil.copy(f, os.path.join(STORE, fn))
            with db.cursor() as c:
                c.execute("INSERT INTO tracks(project_id,name,kind,file_path) VALUES(%s,%s,'stem',%s)",
                          (job["project_id"], f"{src['name']} - {stem}", fn))
                res.append(f"{stem}:{c.lastrowid}")
        with db.cursor() as c: c.execute("UPDATE stem_jobs SET status='done',result_ids=%s WHERE id=%s", (",".join(res), job["id"]))
    except Exception as e:
        with db.cursor() as c: c.execute("UPDATE stem_jobs SET status='failed',message=%s WHERE id=%s", (str(e)[:250], job["id"]))
    finally:
        shutil.rmtree(out, ignore_errors=True)