#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""整谱移调（含调号）+ 旋律清理 + 保存，用于把男声原调抬到女声音区。"""
from __future__ import annotations
import argparse, json, sys, time, urllib.request, urllib.parse, urllib.error

HOST = "http://192.168.31.31:8189"
OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))


def api(path, data=None, timeout=120):
    body = json.dumps(data).encode() if data is not None else None
    req = urllib.request.Request(HOST + path, data=body,
                                headers={"Content-Type": "application/json"} if body else {})
    try:
        with OPENER.open(req, timeout=timeout) as r:
            raw = r.read()
    except urllib.error.HTTPError as e:
        print(f"!! HTTP {e.code}: {e.read()[:1200].decode('utf-8', 'replace')}")
        raise SystemExit(1)
    try:
        return json.loads(raw)
    except Exception:
        return raw


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--abc", required=True)
    p.add_argument("--semitones", type=int, required=True)
    p.add_argument("--scope", default="whole_score",
                   choices=["whole_score", "vocal_only", "instrumental_only", "none"])
    p.add_argument("--cleanup", default="light", choices=["off", "light", "medium", "custom"])
    p.add_argument("--out", required=True)
    a = p.parse_args()

    abc = open(a.abc, encoding="utf-8").read()
    prompt = {
        "1": {"class_type": "YuE2ABCModifier", "inputs": {
            "abc": abc, "tempo_mode": "keep", "bpm": 120.0, "tempo_multiplier": 1.0,
            "transpose_scope": a.scope, "semitones": a.semitones,
            "remove_chords": False, "drop_sections": "",
            "vocal_transpose": 0, "instrumental_transpose": 0, "octave_shift": 0}},
        "2": {"class_type": "YuE2MelodyCleanup", "inputs": {
            "abc": ["1", 0], "preset": a.cleanup,
            "remove_notes_shorter_than_ms": 40.0, "remove_pitch_outliers": True,
            "outlier_semitones": 24}},
        "3": {"class_type": "YuE2SaveABCFile", "inputs": {
            "abc": ["2", 0], "filename": f"YuE2/cover-shift{a.semitones:+d}"}},
    }
    res = api("/prompt", {"prompt": prompt, "client_id": "dsh-shift"})
    if "prompt_id" not in res:
        print("提交失败:", json.dumps(res, ensure_ascii=False)[:1200]); sys.exit(1)
    pid = res["prompt_id"]
    print(f"[提交] 移调 {a.semitones:+d} 半音 scope={a.scope} pid={pid}", flush=True)
    for _ in range(60):
        time.sleep(3)
        h = api(f"/history/{pid}")
        if pid in h:
            rec = h[pid]
            print("status:", rec.get("status", {}).get("status_str"))
            for m in rec.get("status", {}).get("messages") or []:
                if m[0] == "execution_error":
                    print("错误:", json.dumps(m[1], ensure_ascii=False)[:1000])
            outs = rec.get("outputs", {})
            print("产物:", json.dumps(outs, ensure_ascii=False)[:400])
            for nid, o in outs.items():
                for item in (o.get("files") or []):
                    fn = item.get("filename")
                    if not fn:
                        continue
                    url = (f"/view?filename={urllib.parse.quote(fn)}"
                           f"&subfolder={urllib.parse.quote(item.get('subfolder',''))}&type=output")
                    with OPENER.open(HOST + url, timeout=120) as r, open(a.out, "wb") as w:
                        w.write(r.read())
                    print(f"[下载] {a.out}")
            return
    print("超时")


if __name__ == "__main__":
    main()