"""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

import numpy as np
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: godot_rl's ActionSpaceProcessor reports a
    gym.spaces.MultiDiscrete when every per-axis action entry is Discrete
    (see ShipActionCodec/ShipAIController.get_action_space) — nvec
    [5,5,5,5,5,5,2] for rotation xyz, thrust xyz, turbo, in that
    gymnasium-sorted key order. No conversion logic here needs to change for
    that; this class's only functional addition is the truncation-info
    remap below.
    """

    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

    def step(self, action):
        """Remap ShipAIController.get_info()'s "truncated"/"terminal_obs"
        into the keys SB3's on_policy_algorithm looks for
        ("TimeLimit.truncated"/"terminal_observation") so PPO bootstraps
        V(s) through an episode timeout instead of treating every 30s draw
        as a true terminal state.

        Godot_rl's own godot_env.py never sets either key (it returns the
        same `done` array for both term and trunc, "# TODO update API to
        term, trunc") and StableBaselinesGodotEnv.step() only ever returns
        that single collapsed `dones` array to SB3 — so without this, PPO
        has no way to distinguish "episode ended because a goal was scored"
        (a genuine terminal, V(s)=0 is correct) from "episode ended because
        the 30s clock ran out" (an artificial boundary that should be
        bootstrapped through), and was silently treating every draw as the
        former in every curriculum generation to date.
        """
        obs, rewards, dones, infos = super().step(action)
        for info in infos:
            if info.pop("truncated", False):
                info["TimeLimit.truncated"] = True
                info["terminal_observation"] = {"obs": np.array(info.pop("terminal_obs"), dtype=np.float32)}
        return obs, rewards, dones, infos