import json
import os
import signal
import sys
import traceback
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from pathlib import Path

# Force headless Matplotlib to avoid any thread/process GUI issues
os.environ["MPLBACKEND"] = "Agg"

from yolo_webui.config import TrainingConfig
from yolo_webui.trainer import TrainingEvent, TrainingRunner


@contextmanager
def _stop_signal_handlers(runner: TrainingRunner) -> Iterator[None]:
    previous: dict[signal.Signals, signal.Handlers] = {}

    def request_stop(_signum: int, _frame: object) -> None:
        runner.request_stop()

    for signum in (signal.SIGTERM, signal.SIGINT):
        previous[signum] = signal.getsignal(signum)
        signal.signal(signum, request_stop)
    try:
        yield
    finally:
        for signum, handler in previous.items():
            signal.signal(signum, handler)


def main(argv: Sequence[str] | None = None) -> int:
    args = list(sys.argv[1:] if argv is None else argv)
    if not args:
        print(
            "Usage: python -m yolo_webui.subprocess_runner <config_json_path>",
            file=sys.stderr,
        )
        return 1

    config_path = Path(args[0])
    runner = TrainingRunner()
    runner.prepare_run()

    with _stop_signal_handlers(runner):
        # The parent waits for this marker before sending a cooperative signal.
        print("__YOLO_WEBUI_READY__", flush=True)
        try:
            with config_path.open("r", encoding="utf-8") as config_file:
                config_dict = json.load(config_file)
            config = TrainingConfig.from_dict(config_dict)
        except Exception as exc:
            print(f"Error loading config: {exc}", file=sys.stderr)
            return 1

        def handle_event(event: TrainingEvent) -> None:
            event_dict = {
                "kind": event.kind,
                "message": event.message,
                "epoch": event.epoch,
                "total_epochs": event.total_epochs,
            }
            # Print structured JSON event so the parent process can parse it
            print(f"__YOLO_WEBUI_EVENT__:{json.dumps(event_dict)}", flush=True)

        try:
            output_dir = runner.train(config, handle_event)
            if output_dir:
                print(f"__YOLO_WEBUI_RESULT__:{output_dir}", flush=True)
            return 0
        except Exception:
            traceback.print_exc()
            return 1

if __name__ == "__main__":
    raise SystemExit(main())