"""Capture deterministic AutoRemaster UI reference screenshots under Xvfb."""
from __future__ import annotations

import argparse
import subprocess
import time
from pathlib import Path

import numpy as np

from ui_test_support import prepare_headless_ui, reset_module_registry, shutdown_test_app


def wait_for_modules(app, timeout: float = 6.0) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        app.update()
        if app.processing_sidebar.registry._initialized:
            return
        time.sleep(0.005)
    raise RuntimeError("processing modules did not finish loading")


def capture(output: Path, state: str) -> None:
    prepare_headless_ui()
    from ui.app import AutoRemasterApp

    reset_module_registry()
    app = AutoRemasterApp()
    app.geometry("1200x900+0+0")
    wait_for_modules(app)

    if state == "loaded":
        sample_rate = 44100
        t = np.arange(sample_rate * 3, dtype=np.float32) / sample_rate
        left = 0.30 * np.sin(2 * np.pi * 110 * t) + 0.08 * np.sin(2 * np.pi * 2200 * t)
        right = 0.28 * np.sin(2 * np.pi * 113 * t) + 0.07 * np.sin(2 * np.pi * 2600 * t)
        audio = np.column_stack((left, right)).astype(np.float32)
        app.top_toolbar.set_file_label("visual_fixture.wav")
        app.original_section.set_audio(audio, sample_rate)
        app.mastered_section.set_audio(audio * 0.82, sample_rate)
        app.spectrum_display.on_audio_chunk(audio[:2048].mean(axis=1), sample_rate)
        app.update()

    for _ in range(4):
        app.update_idletasks()
        app.update()
        time.sleep(0.03)

    output.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(["scrot", "-a", "0,0,1200,900", str(output)], check=True)
    shutdown_test_app(app)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("output", type=Path)
    parser.add_argument("--state", choices=("initial", "loaded"), default="initial")
    args = parser.parse_args()
    capture(args.output, args.state)