from __future__ import annotations

import argparse
from pathlib import Path

from brittle_star_project import (
    Backend,
    BrittleStarEnv,
    BrittleStarEnvFactory,
    SimulationConfig,
    simulate_policy,
)
from brittle_star_project.environment import from_json
from brittle_star_project.rl import RLModel  # imports concrete models via rl.__init__
from brittle_star_project.rl.base import get_rl_model_registry

MODEL_BY_NAME = get_rl_model_registry()
MODEL_OPTIONS = sorted(MODEL_BY_NAME)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Simulate a trained policy in the MuJoCo viewer.")
    p.add_argument(
        "--model",
        type=str,
        default=None,
        help="Path to a saved model artifact. If omitted, a model is created from --model-type.",
    )
    p.add_argument(
        "--model-type",
        choices=MODEL_OPTIONS,
        default="random",
        help="Which model class to instantiate when --model is omitted.",
    )
    p.add_argument(
        "--backend",
        choices=[b for b in Backend],
        default=Backend.MJX,
    )
    p.add_argument("--seed", type=int, default=None)
    return p.parse_args()


def main() -> None:
    args = parse_args()

    morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json")

    # ======= ENVIRONMENT SETUP =======

    backend = args.backend

    factory = BrittleStarEnvFactory()
    raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
    env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg)

    seed_for_env = int(args.seed) if args.seed is not None else 0
    state = env.reset(seed=seed_for_env)

    # ======= MODEL SETUP =======

    # Extract the number of actuators (nu) from the environment's model, so we can pass it to the
    # policy/model.
    nu = int(state.mj_model.nu)

    if args.model is not None:
        model_path = Path(args.model)
        policy = RLModel.load(model_path)
        if hasattr(policy, "nu"):
            policy.nu = nu
    else:
        model_cls = MODEL_BY_NAME[str(args.model_type)]
        policy = model_cls(seed=seed_for_env)
        if hasattr(policy, "nu"):
            policy.nu = nu

    # If the policy/model has a `seed` attribute, use the provided seed (or default) to reset it.
    default_seed = int(getattr(policy, "seed", seed_for_env))
    if args.seed is not None and hasattr(policy, "reset"):
        policy.reset(int(args.seed))

    # ======= SIMULATION =======

    rollout_cfg = SimulationConfig(
        realtime=True,
        seed=int(args.seed) if args.seed is not None else default_seed,
    )

    simulate_policy(policy, rollout_cfg, state)

    env.close()


if __name__ == "__main__":
    main()