"""Gymnasium environment that wraps the GameCli C# binary as a subprocess.

Speaks the line-delimited JSON protocol defined in:
  docs/superpowers/specs/2026-07-17-cursor-following-lander-design.md

Each episode picks a random target in [0,1]^2 at reset time and holds it fixed.
Never terminates — episodes end via truncation at max_episode_steps.
"""
from __future__ import annotations

import json
import subprocess
from pathlib import Path
from typing import Any, Optional

import gymnasium as gym
import numpy as np
from gymnasium import spaces


class LanderCliEnv(gym.Env):
    """One CLI subprocess per env instance. Not thread-safe; safe under SubprocVecEnv."""

    metadata = {"render_modes": []}

    def __init__(
        self,
        cli_binary: str,
        max_episode_steps: int = 1000,
    ) -> None:
        super().__init__()
        self._cli_binary = cli_binary
        self._max_episode_steps = max_episode_steps

        self._proc: Optional[subprocess.Popen] = None
        self._target: tuple[float, float] = (0.5, 0.5)
        self._step_count: int = 0

        self._start_process()
        self._read_init()

        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32
        )
        self.action_space = spaces.Discrete(4)

    # -------- subprocess lifecycle --------

    def _start_process(self) -> None:
        self._proc = subprocess.Popen(
            [self._cli_binary],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,  # line-buffered
        )

    def _read_init(self) -> None:
        assert self._proc is not None and self._proc.stdout is not None
        line = self._proc.stdout.readline()
        if not line:
            raise RuntimeError("CLI died before emitting init handshake")
        msg = json.loads(line)
        if "init" not in msg:
            raise RuntimeError(f"expected init handshake, got: {line!r}")

    def _send(self, obj: dict[str, Any]) -> None:
        assert self._proc is not None and self._proc.stdin is not None
        self._proc.stdin.write(json.dumps(obj) + "\n")
        self._proc.stdin.flush()

    def _recv(self) -> dict[str, Any]:
        assert self._proc is not None and self._proc.stdout is not None
        line = self._proc.stdout.readline()
        if not line:
            raise RuntimeError("CLI closed stdout unexpectedly")
        return json.loads(line)

    # -------- gym.Env API --------

    def reset(
        self, *, seed: Optional[int] = None, options: Optional[dict] = None
    ) -> tuple[np.ndarray, dict]:
        super().reset(seed=seed)
        # Sample a fresh target each episode; the Gymnasium-provided
        # np_random is deterministic given `seed`.
        tx = float(self.np_random.uniform(0.05, 0.95))
        ty = float(self.np_random.uniform(0.05, 0.95))
        self._target = (tx, ty)
        self._step_count = 0

        cmd: dict[str, Any] = {"cmd": "reset"}
        if seed is not None:
            cmd["seed"] = int(seed)
        self._send(cmd)
        msg = self._recv()
        obs = np.asarray(msg["obs"], dtype=np.float32)
        # Overwrite the CLI's dummy (0.5, 0.5) target displacement with the real one.
        # dx = target.X - ship.X, dy = target.Y - ship.Y.
        obs[0] = tx - msg["state"]["x"]
        obs[1] = ty - msg["state"]["y"]
        return obs, {}

    def step(self, action: int) -> tuple[np.ndarray, float, bool, bool, dict]:
        self._send({"action": int(action), "target": list(self._target)})
        msg = self._recv()
        obs = np.asarray(msg["obs"], dtype=np.float32)
        reward = float(msg["reward"])
        self._step_count += 1
        truncated = self._step_count >= self._max_episode_steps
        terminated = False  # continuous hover task
        return obs, reward, terminated, truncated, {}

    def close(self) -> None:
        if self._proc is None:
            return
        try:
            if self._proc.stdin is not None:
                self._proc.stdin.close()
        except (BrokenPipeError, OSError):
            pass
        try:
            self._proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            self._proc.kill()
            self._proc.wait(timeout=1)
        self._proc = None

    def __del__(self) -> None:
        try:
            self.close()
        except Exception:
            pass