import gymnasium as gym
import numpy as np
import json
import time
from .client.elden_client import EldenClient
from .rewards import RewardFunction, ScoreDeltaReward


class EldenGymEnv(gym.Env):
    """
    Elden Ring Gymnasium environment with non-blocking frame streaming.

    Uses pysiphon's frame streaming for efficient polling-based observations.

    Args:
        scenario_name (str): Boss scenario name
        keybinds_filepath (str): Path to keybinds JSON file
        siphon_config_filepath (str): Path to siphon TOML config
        memory_attributes (list[str]): List of memory attribute names to include in observation.
            Default: ["HeroHp", "HeroMaxHp", "NpcHp", "NpcMaxHp", "HeroAnimId", "NpcAnimId"]
        host (str): Siphon server host. Default: 'localhost:50051'
        reward_function (RewardFunction): Custom reward function
        frame_format (str): Frame format for streaming ('jpeg' or 'raw'). Default: 'jpeg'
        frame_quality (int): JPEG quality 1-100. Default: 85
        max_steps (int): Maximum steps per episode. Default: None
        launch_game (bool): Whether to launch the game automatically. Default: True
            Set to False if game is already running
        save_file_name (str, optional): Name of backup save file to copy during reset
            (e.g., "margit_checkpoint.sl2"). If None, no save file copying occurs.
        save_file_dir (str, optional): Directory containing save files. Required if
            save_file_name is provided. Typically: %APPDATA%/EldenRing/<steam_id>/
    """

    def __init__(
        self,
        scenario_name,
        keybinds_filepath,
        siphon_config_filepath,
        memory_attributes=None,
        host="localhost:50051",
        reward_function=None,
        frame_format="jpeg",
        frame_quality=85,
        max_steps=None,
        launch_game=True,
        save_file_name=None,
        save_file_dir=None,
    ):
        super().__init__()

        self.scenario_name = scenario_name
        self.client = EldenClient(host)
        self.keybinds_filepath = keybinds_filepath
        self.siphon_config_filepath = siphon_config_filepath
        self.step_count = 0
        self.max_steps = max_steps
        self.frame_format = frame_format
        self.frame_quality = frame_quality
        self.save_file_name = save_file_name
        self.save_file_dir = save_file_dir

        # Memory attributes to poll (configurable, not hardcoded)
        self.memory_attributes = memory_attributes or [
            "HeroHp",
            "HeroMaxHp",
            "NpcHp",
            "NpcMaxHp",
            "HeroAnimId",
            "NpcAnimId",
        ]

        # Load keybinds
        with open(self.keybinds_filepath, "r") as f:
            keybinds_data = json.load(f)
            self.keybinds = keybinds_data["keybinds"]

        # Create action space (multi-binary for all keys)
        self.action_keys = list(self.keybinds.keys())
        self.action_space = gym.spaces.MultiBinary(len(self.action_keys))

        # Track current key states for toggling
        self._key_states = {key: False for key in self.action_keys}

        # Frame stream handle
        self._stream_handle = None

        # Reward function
        self.reward_function = reward_function or ScoreDeltaReward(
            score_key="player_hp"
        )
        if not isinstance(self.reward_function, RewardFunction):
            raise TypeError("reward_fn must inherit from RewardFunction")

        # State tracking
        self._prev_info = None

        # Initialize game and siphon
        if launch_game:
            print("Launching game...")
            self.client.launch_game()
            time.sleep(20)  # Wait for game to launch
            self.client.enter_game()
        else:
            print("Skipping game launch (launch_game=False)")

        print("Initializing Siphon...")
        self.client.load_config_from_file(self.siphon_config_filepath, wait_time=2)
        time.sleep(2)

        # Verify server is ready
        print("Checking server status...")
        status = self.client.get_server_status()
        print(f"Server status: {status}")

        if not status.get("memory_initialized", False):
            raise RuntimeError("Memory subsystem not initialized!")
        if not status.get("capture_initialized", False):
            raise RuntimeError("Capture subsystem not initialized!")

        print("Starting frame stream...")
        self._stream_handle = self.client.start_frame_stream(
            format=self.frame_format, quality=self.frame_quality
        )

        # Setup observation space (will be defined after first observation)
        self.observation_space = None

    def _poll_observation(self):
        """
        Poll for latest frame and memory attributes.

        Returns:
            dict: Observation with 'frame' and memory attributes
        """
        # Poll latest frame (non-blocking)
        frame_data = self.client.get_latest_frame(self._stream_handle)

        # If no new frame available, wait briefly and retry
        if frame_data is None:
            time.sleep(0.005)
            frame_data = self.client.get_latest_frame(self._stream_handle)

        # Decode frame from protobuf FrameData object
        if frame_data is not None:
            import cv2
            import numpy as np

            # Extract JPEG bytes from protobuf
            jpeg_bytes = frame_data.data

            if jpeg_bytes and len(jpeg_bytes) > 0:
                # Decode JPEG bytes to numpy array
                frame_array = np.frombuffer(jpeg_bytes, dtype=np.uint8)
                frame = cv2.imdecode(frame_array, cv2.IMREAD_COLOR)
            else:
                # Fallback: create black frame if no data
                frame = np.zeros(
                    (frame_data.height, frame_data.width, 3), dtype=np.uint8
                )
        else:
            # No frame available, create black placeholder
            frame = np.zeros((2160, 3840, 3), dtype=np.uint8)

        # Get memory attributes
        memory_data = {}
        for attr_name in self.memory_attributes:
            try:
                response = self.client.get_attribute(attr_name)
                # pysiphon returns dict with 'value' key
                if isinstance(response, dict):
                    value = response.get("value", 0)
                else:
                    value = response
                memory_data[attr_name] = value
            except Exception as e:
                print(f"Warning: Could not read attribute {attr_name}: {e}")
                print(
                    "  This might mean the game isn't fully loaded yet or the attribute doesn't exist."
                )
                memory_data[attr_name] = 0

        # Combine into observation
        obs = {"frame": frame, **memory_data}

        return obs

    def _toggle_keys(self, action):
        """
        Toggle keys based on multi-binary action and current key states.

        Args:
            action: Multi-binary array indicating desired key states
        """
        for i, desired_state in enumerate(action):
            key = self.action_keys[i]
            current_state = self._key_states[key]
            new_state = bool(desired_state)

            # Only toggle if state changed
            if new_state != current_state:
                self.client.input_key_toggle(key, new_state)
                self._key_states[key] = new_state

    def _release_all_keys(self):
        """Release all currently pressed keys."""
        for key, is_pressed in self._key_states.items():
            if is_pressed:
                self.client.input_key_toggle(key, False)
                self._key_states[key] = False

    def reset(self, seed=None, options=None):
        """Reset environment - start new episode."""
        super().reset(seed=seed)

        # Release all keys from previous episode
        self._release_all_keys()

        # Reset game state
        self.client.quit_to_title()

        # Copy save file if configured
        if self.save_file_name and self.save_file_dir:
            print(f"Copying save file: {self.save_file_name}")
            self.client.copy_save_file(self.save_file_name, self.save_file_dir)

        self.client.enter_menu()
        self.client.start_scenario(self.scenario_name)

        # Reset tracking
        self.step_count = 0
        self._prev_info = None

        # Get initial observation
        obs = self._poll_observation()

        # Define observation space on first reset if not already defined
        if self.observation_space is None:
            self.observation_space = gym.spaces.Dict(
                {
                    "frame": gym.spaces.Box(
                        low=0,
                        high=255,
                        shape=obs["frame"].shape,
                        dtype=np.uint8,
                    ),
                    **{
                        attr: gym.spaces.Box(
                            low=-np.inf, high=np.inf, shape=(), dtype=np.float32
                        )
                        for attr in self.memory_attributes
                    },
                }
            )

        info = self._get_info(obs)
        self._prev_info = info.copy()

        return obs, info

    def step(self, action):
        """
        Execute one step with key toggling.

        Args:
            action: Multi-binary array [0/1] for each key in self.action_keys

        Returns:
            tuple: (observation, reward, terminated, truncated, info)
        """
        # Toggle keys based on action
        self._toggle_keys(action)

        # Brief wait for game to process input
        time.sleep(0.016)  # ~1 frame at 60fps

        # Poll observation
        obs = self._poll_observation()
        info = self._get_info(obs)

        # Calculate reward
        reward = self.reward_function.calculate(obs, info, self._prev_info)

        # Check termination
        terminated = self.reward_function.is_done(obs, info)
        truncated = (
            self.step_count >= self.max_steps if self.max_steps is not None else False
        )

        # Update tracking
        self.step_count += 1
        self._prev_info = info.copy()

        return obs, reward, terminated, truncated, info

    def _get_info(self, obs):
        """
        Extract info dict from observation.

        Args:
            obs: Observation dict

        Returns:
            dict: Info with normalized/processed values
        """
        info = {}

        # Add normalized HP values if available
        if "HeroHp" in obs and "HeroMaxHp" in obs:
            info["player_hp_normalized"] = (
                obs["HeroHp"] / obs["HeroMaxHp"] if obs["HeroMaxHp"] > 0 else 0
            )

        if "NpcHp" in obs and "NpcMaxHp" in obs:
            info["boss_hp_normalized"] = (
                obs["NpcHp"] / obs["NpcMaxHp"] if obs["NpcMaxHp"] > 0 else 0
            )

        # Add animation IDs
        if "HeroAnimId" in obs:
            info["player_animation"] = obs["HeroAnimId"]

        if "NpcAnimId" in obs:
            info["boss_animation"] = obs["NpcAnimId"]

        return info

    def close(self):
        """Close environment and clean up resources."""
        # Stop frame stream
        if self._stream_handle is not None:
            self.client.stop_frame_stream(self._stream_handle)
            self._stream_handle = None

        # Release all keys
        self._release_all_keys()

        # Close client
        self.client.close()

    def render(self):
        """Render is handled by the game itself."""
        pass
