"""Live ViZDoom/Freedoom runtime for ARGOS-MUKHTAR Space.

The Space owns the game world. A remote controller may only submit actions
when MUKHTAR_BRIDGE_KEY is configured and the caller supplies the same key.
No original DOOM assets are bundled; ViZDoom uses its free Freedoom assets.
"""
from __future__ import annotations

import os
import threading
import time
from typing import Any

import numpy as np

ACTION_NAMES = ("NOOP", "FORWARD", "BACK", "TURN_LEFT", "TURN_RIGHT", "FIRE")
_ACTIONS = {
    "NOOP":        [0, 0, 0, 0, 0],
    "FORWARD":     [1, 0, 0, 0, 0],
    "BACK":        [0, 1, 0, 0, 0],
    "TURN_LEFT":   [0, 0, 1, 0, 0],
    "TURN_RIGHT":  [0, 0, 0, 1, 0],
    "FIRE":        [0, 0, 0, 0, 1],
}
_LOCK = threading.RLock()
class DoomRuntime:
    def __init__(self) -> None:
        self.game = None
        self.seed = 0
        self.episode = 0
        self.step_count = 0
        self.total_reward = 0.0
        self.last_reward = 0.0
        self.last_action = "NOOP"
        self.started_at = None

    def _build(self, seed: int):
        import vizdoom as vzd
        game = vzd.DoomGame()
        game.set_screen_resolution(vzd.ScreenResolution.RES_320X240)
        game.set_screen_format(vzd.ScreenFormat.RGB24)
        game.set_window_visible(False)
        game.set_mode(vzd.Mode.PLAYER)
        game.set_doom_map("map01")
        game.set_doom_skill(3)
        game.set_seed(int(seed))
        game.set_available_buttons([
            vzd.Button.MOVE_FORWARD,
            vzd.Button.MOVE_BACKWARD,
            vzd.Button.TURN_LEFT,
            vzd.Button.TURN_RIGHT,
            vzd.Button.ATTACK,
        ])
        game.set_available_game_variables([
            vzd.GameVariable.HEALTH,
            vzd.GameVariable.ARMOR,
            vzd.GameVariable.SELECTED_WEAPON_AMMO,
            vzd.GameVariable.KILLCOUNT,
            vzd.GameVariable.DEATHCOUNT,
        ])
        game.set_living_reward(-0.001)
        game.init()
        return game

    def start(self, seed: int = 42):
        with _LOCK:
            if self.game is not None:
                try:
                    self.game.close()
                except Exception:
                    pass
            self.seed = int(seed)
            self.game = self._build(self.seed)
            self.game.new_episode()
            self.episode += 1
            self.step_count = 0
            self.total_reward = 0.0
            self.last_reward = 0.0
            self.last_action = "NOOP"
            self.started_at = time.time()
            return self.snapshot()
    def _ensure(self):
        if self.game is None:
            return self.start(self.seed or 42)
        if self.game.is_episode_finished():
            self.game.new_episode()
            self.episode += 1
            self.step_count = 0
            self.total_reward = 0.0
        return None

    def step(self, action: str, tics: int = 4):
        name = str(action or "NOOP").upper()
        if name not in _ACTIONS:
            raise ValueError(f"unknown Doom action: {name}")
        frames = max(1, min(int(tics), 12))
        with _LOCK:
            self._ensure()
            reward = float(self.game.make_action(_ACTIONS[name], frames))
            self.step_count += 1
            self.total_reward += reward
            self.last_reward = reward
            self.last_action = name
            if self.game.is_episode_finished():
                return self.snapshot(finished=True)
            return self.snapshot()
    def _var(self, name: str, default: float = 0.0) -> float:
        import vizdoom as vzd
        try:
            return float(self.game.get_game_variable(getattr(vzd.GameVariable, name)))
        except Exception:
            return float(default)

    def snapshot(self, finished: bool = False):
        with _LOCK:
            if self.game is None:
                return None, self.telemetry(state="not_started")
            state = None if self.game.is_episode_finished() else self.game.get_state()
            frame = None if state is None else np.asarray(state.screen_buffer).copy()
            info = self.telemetry(state="finished" if finished else "running")
            return frame, info

    def telemetry(self, state: str = "running") -> dict[str, Any]:
        return {
            "engine": "ViZDoom",
            "assets": "Freedoom (free content)",
            "state": state,
            "seed": self.seed,
            "episode": self.episode,
            "step": self.step_count,
            "last_action": self.last_action,
            "last_reward": round(self.last_reward, 6),
            "total_reward": round(self.total_reward, 6),
            "health": self._var("HEALTH"),
            "armor": self._var("ARMOR"),
            "ammo": self._var("SELECTED_WEAPON_AMMO"),
            "kills": self._var("KILLCOUNT"),
            "deaths": self._var("DEATHCOUNT"),
            "uptime_s": round(time.time() - self.started_at, 2)
            if self.started_at else 0.0,
            "remote_bridge": bridge_status(),
        }


_RUNTIME = DoomRuntime()


def bridge_status() -> dict[str, Any]:
    configured = bool(os.getenv("MUKHTAR_CONTROL_CODE", "").strip())
    return {
        "configured": configured,
        "remote_write": "enabled" if configured else "disabled",
        "policy": "controller-code protected; observation is read-only",
    }
def doom_start(seed: int = 42):
    return _RUNTIME.start(seed)


def doom_step(action: str, tics: int = 4):
    return _RUNTIME.step(action, tics)


def doom_snapshot():
    return _RUNTIME.snapshot()


def doom_observe():
    return _RUNTIME.snapshot()


def doom_remote_step(action: str, tics: int, control_code: str):
    expected = os.getenv("MUKHTAR_CONTROL_CODE", "").strip()
    supplied = str(control_code or "")
    frame, info = _RUNTIME.snapshot()
    if not expected:
        info = dict(info)
        info["remote_error"] = "control_not_configured"
        return frame, info
    import secrets
    if not secrets.compare_digest(expected, supplied):
        info = dict(info)
        info["remote_error"] = "unauthorized"
        return frame, info
    return _RUNTIME.step(action, tics)