#!/usr/bin/env python3
"""资源保护自检：小分段生效，重任务不会并行抢机器。"""
import os
import sys
import tempfile
import threading
import time

if len(sys.argv) > 1:
    arg = sys.argv[1]
    if arg in ("-h", "--help"):
        print(__doc__.strip())
        sys.exit(0)
    if arg in ("-v", "--version"):
        print("fenxuan resource guard 1.5.0")
        sys.exit(0)
    print(f"unknown arguments: {sys.argv[1:]}", file=sys.stderr)
    sys.exit(2)

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "server"))
os.environ["FENXUAN_DATA"] = tempfile.mkdtemp(prefix="fenxuan-guard-")

import app as server  # noqa: E402
import separate  # noqa: E402


def main():
    assert server._track_filename("像风/一样", "钢琴", ext="wav") == \
        "像风_一样_钢琴.wav"
    assert server._passes_quality(-10, 0.5)
    assert not server._passes_quality(-30, 0.5)
    assert not server._passes_quality(-10, 0.01)
    assert server.OUTPUT_FORMATS["both"] == ("mp3", "wav")
    assert set(server.PROFILES) == {"vocal", "stems6"}
    assert separate.STEMS_6S == (
        "vocals", "drums", "bass", "other", "guitar", "piano")

    os.environ.pop("FENXUAN_ROFORMER_SEGMENT", None)
    assert separate._roformer_params() == {
        "segment_size": 512, "batch_size": 1, "overlap": 4,
        "override_model_segment_size": True}
    os.environ["FENXUAN_ROFORMER_SEGMENT"] = "32"
    assert separate._roformer_params()["segment_size"] == 64
    assert separate._roformer_params()["overlap"] == 1
    os.environ["FENXUAN_ROFORMER_SEGMENT"] = "999"
    assert separate._roformer_params()["segment_size"] == 512
    assert separate._roformer_params()["overlap"] == 4
    os.environ["FENXUAN_ROFORMER_SEGMENT"] = "bad"
    assert separate._roformer_segment() == 512

    active = 0
    peak = 0
    guard = threading.Lock()

    def work(job_id):
        nonlocal active, peak
        with guard:
            active += 1
            peak = max(peak, active)
        time.sleep(0.08)
        with guard:
            active -= 1

    for job_id in ("a", "b"):
        server.JOBS[job_id] = {"stage": "decode", "progress": 0}
    threads = [threading.Thread(target=server._run_one, args=(work, job_id))
               for job_id in ("a", "b")]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()
    assert peak == 1
    print("RESOURCE_GUARD: PASS")


if __name__ == "__main__":
    main()
