#!/usr/bin/env python3

"""Run a bounded set of simulation variants concurrently."""

import argparse
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from importlib import import_module
import json
import os
from pathlib import Path
from typing import Any, cast


JsonObject = dict[str, Any]
BuildRequest = Callable[..., JsonObject]
RunSimulation = Callable[..., Path]
VARIANTS = {
    "8-fins-40-w": {"num_fins": 8, "power": 40.0},
    "10-fins-40-w": {"num_fins": 10, "power": 40.0},
    "8-fins-50-w": {"num_fins": 8, "power": 50.0},
    "10-fins-50-w": {"num_fins": 10, "power": 50.0},
}


def load_local_helpers() -> tuple[BuildRequest, RunSimulation]:
    """Load the request builder and runner from the reader's simulation script."""
    # The reader adds study parameters to build_request in the script beside this driver.
    runner = import_module("run_simulation")
    build_request = getattr(runner, "build_request", None)
    run_simulation = getattr(runner, "run_simulation", None)
    if not callable(build_request) or not callable(run_simulation):
        raise RuntimeError("run_simulation.py must define both build_request and run_simulation")
    return cast(BuildRequest, build_request), cast(RunSimulation, run_simulation)


def parse_arguments() -> argparse.Namespace:
    """Read the concurrency limit and study directory."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--max-workers",
        type=int,
        required=True,
        help="maximum simultaneous jobs; do not exceed your API dashboard allowance",
    )
    parser.add_argument("--output", type=Path, default=Path("results/design-study"))
    arguments = parser.parse_args()
    if arguments.max_workers < 1:
        parser.error("--max-workers must be at least 1")
    return arguments


def save_requests(request_directory: Path, build_request: BuildRequest) -> dict[str, JsonObject]:
    """Construct and save every request before starting any simulations."""
    request_directory.mkdir(parents=True, exist_ok=True)
    requests: dict[str, JsonObject] = {}
    for name, parameters in VARIANTS.items():
        request = build_request(**parameters)
        (request_directory / f"{name}.json").write_text(
            json.dumps(request, indent=2) + "\n",
            encoding="utf-8",
        )
        requests[name] = request
    return requests


def main() -> None:
    """Prepare every request, then run them within the selected worker limit."""
    arguments = parse_arguments()
    api_key = os.environ.get("VANELLUS_API_KEY")
    if not api_key:
        raise SystemExit("Set VANELLUS_API_KEY before running this script")

    build_request, run_simulation = load_local_helpers()
    requests = save_requests(arguments.output / "requests", build_request)
    result_directory = arguments.output / "results"
    futures: dict[Future[Path], str] = {}
    with ThreadPoolExecutor(max_workers=arguments.max_workers) as executor:
        for name, request in requests.items():
            future = executor.submit(
                run_simulation,
                request=request,
                api_key=api_key,
                output_root=result_directory,
                output_name=name,
            )
            futures[future] = name

        failed = False
        for future in as_completed(futures):
            name = futures[future]
            try:
                print(f"{name}: results saved to {future.result()}")
            except Exception as error:
                failed = True
                print(f"{name} failed: {error}")

    if failed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()