"""Godot RL Agents environment wrappers that run Cosmic Clash from source.

Stock GodotEnv expects an *exported* game executable and rewrites its path
per-platform. These subclasses launch the project straight from the repo with
a Godot binary instead (no export step), pointing it at the training scene.
Each Godot instance contributes two agents (one ship per team) that share the
learning policy: self-play by construction.
"""

import pathlib
import subprocess

from godot_rl.core.godot_env import GodotEnv
from godot_rl.wrappers.stable_baselines_wrapper import StableBaselinesGodotEnv

REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
GAME_DIR = REPO_ROOT / "Game"
TRAINING_SCENE = "res://scenes/training.tscn"


class CosmicClashEnv(GodotEnv):
    """GodotEnv that launches either the project from source or an exported binary.

    Source mode (default): `godot --path Game res://scenes/training.tscn` — the
    positional scene argument overrides the project's normal main scene.
    Exported mode (`exported=True`): `env_path` is a pre-built game executable
    (see training/export_linux.sh, "Linux Training" preset) — no `--path` and
    no scene override needed or possible: official Godot export templates have
    path/scene overrides compiled out (`--scene`/a positional scene argument
    hard-aborts with "compiled without support for path overrides"), so the
    binary instead boots straight into training.tscn on its own via
    project.godot's `run/main_scene.training` feature-tag override, activated
    by that preset's `custom_features="training"`.
    """

    def __init__(self, *args, exported: bool = False, **kwargs):
        self.exported = exported
        super().__init__(*args, **kwargs)

    # env_path is a Godot binary (or, in exported mode, a game executable we
    # built ourselves), not a stock godot_rl exported-project path: skip the
    # suffix and platform checks stock GodotEnv applies to those.
    def _set_platform_suffix(self, env_path: str) -> str:
        return env_path

    def check_platform(self, filename: str):
        pass

    def _launch_env(self, env_path, port, show_window, framerate, seed, action_repeat, speedup, **kwargs):
        # sync.gd reads --key=value pairs from the raw command line; they must
        # NOT go after a `--` separator or OS.get_cmdline_args() drops them.
        cmd = [env_path]
        if not self.exported:
            cmd += ["--path", str(GAME_DIR), TRAINING_SCENE]
        cmd += [
            f"--port={port}",
            f"--env_seed={seed}",
        ]
        if not show_window:
            cmd += ["--headless", "--disable-render-loop"]
        if framerate is not None:
            cmd += ["--fixed-fps", str(framerate)]
        if action_repeat is not None:
            cmd.append(f"--action_repeat={action_repeat}")
        if speedup is not None:
            cmd.append(f"--speedup={speedup}")
        for key, value in kwargs.items():
            cmd.append(f"--{key}={value}")
        self.proc = subprocess.Popen(cmd, start_new_session=True)


class CosmicClashVecEnv(StableBaselinesGodotEnv):
    """SB3 VecEnv over N parallel CosmicClashEnv instances.

    convert_action_space=True flattens the env's (Box(6), Discrete(2)) action
    space into a single Box(7): thrust xyz, rotation xyz, turbo (>0 means on).
    """

    def __init__(self, godot_bin: str, n_parallel: int = 1, seed: int = 0, port: int = GodotEnv.DEFAULT_PORT, **kwargs):
        self.envs = [
            CosmicClashEnv(
                env_path=godot_bin,
                convert_action_space=True,
                port=port + p,
                seed=seed + p,
                **kwargs,
            )
            for p in range(n_parallel)
        ]
        self.n_parallel = n_parallel
        self._check_valid_action_space()
        self.results = None