from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")

from hamworld.runtime import ExperimentLogger, apply_overrides, expand_task_runs, load_config, make_env, module_summary
from hamworld.world_model import CanonicalDynamicsWorldModel


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Visualize the canonical dynamics HaM-World architecture.")
    parser.add_argument(
        "--config",
        type=str,
        default=str(Path(__file__).resolve().parent / "configs" / "default.yaml"),
        help="Path to a YAML config.",
    )
    parser.add_argument("--task", default=None, type=str, help="Optional task id to run.")
    parser.add_argument("--override", action="append", default=[], help="Dotted config override.")
    return parser.parse_args()


def run_from_config(config: dict) -> None:
    task = config["task"]
    env, spec = make_env(task, int(config["experiment"]["seed"]))
    del env
    logger = ExperimentLogger(
        output_root=config["experiment"]["output_dir"],
        algorithm=config["experiment"]["algorithm"],
        task_id=f"{task['id']}_viz",
        seed=int(config["experiment"]["seed"]),
    )
    model = CanonicalDynamicsWorldModel(config, spec.observation_shape[0], spec.action_shape[0])
    summary = module_summary(model)
    logger.info(summary)
    logger.save_json(
        "architecture.json",
        {
            "algorithm": config["experiment"]["algorithm"],
            "task": task["id"],
            "summary": summary,
        },
    )


def main(config: dict | None = None) -> None:
    if config is not None:
        run_from_config(config)
        return

    args = parse_args()
    resolved_config = apply_overrides(load_config(Path(args.config)), args.override)
    for run_config in expand_task_runs(resolved_config, task_filter=args.task):
        run_from_config(run_config)


if __name__ == "__main__":
    main()