import _ra_wrapper
import numpy as np
import os
from dataclasses import dataclass
from enum import IntEnum

class RetroButton(IntEnum):
    """RetroArch button IDs matching libretro RETRO_DEVICE_ID_JOYPAD_*"""
    B = 0
    Y = 1
    SELECT = 2
    START = 3
    UP = 4
    DOWN = 5
    LEFT = 6
    RIGHT = 7
    A = 8
    X = 9
    L = 10
    R = 11
    L2 = 12
    R2 = 13
    L3 = 14
    R3 = 15

class RetroAnalogStick(IntEnum):
    """Analog stick selection"""
    LEFT = 0
    RIGHT = 1

class RetroAnalogAxis(IntEnum):
    """Analog axis selection"""
    X = 0
    Y = 1

@dataclass
class SystemInfo:
    library_name: str
    library_version: str
    valid_extensions: list[str]
    need_fullpath: bool

    @classmethod
    def from_dict(cls, info: dict) -> 'SystemInfo':
        return cls(
            library_name=info['library_name'],
            library_version=info['library_version'],
            valid_extensions=info['valid_extensions'].split('|'),
            need_fullpath=info['need_fullpath']
        )

@dataclass
class AVInfo:
    fps: float
    sample_rate: float
    base_width: int
    base_height: int
    max_width: int
    max_height: int
    aspect_ratio: float

    @classmethod
    def from_dict(cls, info: dict) -> 'AVInfo':
        return cls(**info)

@dataclass
class VideoFrame:
    data: np.ndarray  # Raw pixel data (pointer / bytes)
    width: int
    height: int
    pitch: int
    format: int

    @classmethod
    def from_dict(cls, frame: dict) -> 'VideoFrame':
        return cls(**frame)

    def to_rgb(self) -> np.ndarray:
        # Allocate output array once
        out = np.empty((self.height, self.width, 3), dtype=np.uint8)

        if self.format == 1:  # XRGB8888
            raw = np.frombuffer(self.data, dtype=np.uint32).reshape(self.height, self.pitch // 4)
            raw = raw[:, :self.width]  # just view, no copy
            out[..., 0] = (raw >> 16) & 0xFF  # R
            out[..., 1] = (raw >> 8) & 0xFF   # G
            out[..., 2] = raw & 0xFF          # B

        elif self.format == 0:  # 0RGB1555
            raw = np.frombuffer(self.data, dtype=np.uint16).reshape(self.height, self.pitch // 2)
            raw = raw[:, :self.width]
            out[..., 0] = ((raw >> 10) & 0x1F) << 3
            out[..., 1] = ((raw >> 5) & 0x1F) << 3
            out[..., 2] = (raw & 0x1F) << 3

        elif self.format == 2:  # RGB565
            raw = np.frombuffer(self.data, dtype=np.uint16).reshape(self.height, self.pitch // 2)
            raw = raw[:, :self.width]
            out[..., 0] = ((raw >> 11) & 0x1F) << 3
            out[..., 1] = ((raw >> 5) & 0x3F) << 2
            out[..., 2] = (raw & 0x1F) << 3

        else:
            raise ValueError(f"Unsupported video format: {self.format}")

        return out



@dataclass
class AudioFrame:
    data: np.ndarray  # Stereo int16 samples
    frames: int
    sample_rate: float

    @classmethod
    def from_dict(cls, audio: dict) -> 'AudioFrame':
        return cls(**audio)

class Emulator:
    def __init__(self, core_path: str):
        self.core_path = os.path.abspath(core_path)
        
        if not os.path.isfile(core_path):
            raise FileNotFoundError(f"Core not found: {core_path}")
        
        _ra_wrapper.init_core(self.core_path)
        self._system_info = SystemInfo.from_dict(_ra_wrapper.get_system_info())
        self._av_info = None
        self._game_loaded = False
    
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        try: 
            _ra_wrapper.shutdown()
        except:
            pass
    
    @property
    def system_info(self) -> SystemInfo:
        return self._system_info
    
    @property
    def av_info(self) -> AVInfo:
        if not self._game_loaded or self._av_info is None:
            raise RuntimeError("No game loaded")
        return self._av_info

    
    @property
    def is_game_loaded(self) -> bool:
        return self._game_loaded

    def load_game(self, rom_path: str):
        rom_path = os.path.abspath(rom_path)
        
        if not os.path.isfile(rom_path):
            raise FileNotFoundError(f"ROM not found: {rom_path}")
        
        _ra_wrapper.load_game(rom_path)
        self._av_info = AVInfo.from_dict(_ra_wrapper.get_av_info())
        self._game_loaded = True
    
    def unload_game(self):
        if self._game_loaded:
            _ra_wrapper.unload_game()
            self._game_loaded = False
            self._av_info = None
    
    def reset(self):
        if not self._game_loaded:
            raise RuntimeError("No game loaded")
        _ra_wrapper.reset()
    
    def step(self):
        if not self._game_loaded:
            raise RuntimeError("No game loaded")
        _ra_wrapper.step()
    
    def get_video_frame(self) -> VideoFrame:
        return VideoFrame.from_dict(_ra_wrapper.get_video_frame())
    
    def get_audio_frame(self) -> AudioFrame:
        return AudioFrame.from_dict(_ra_wrapper.get_audio_frame())

    def set_button(self, port: int, button: int | RetroButton, pressed: bool):
        """Set button state
        
        Args:
            port: Controller port (0-3)
            button: Button ID (0-15) or RetroButton enum
            pressed: True if pressed, False if released
        
        Example:
            emu.set_button(0, RetroButton.A, True)
            emu.set_button(0, 8, True)  # Also valid
        """
        _ra_wrapper.set_button(port, int(button), pressed)

    def set_analog(self, port: int, stick: int | RetroAnalogStick, axis: int | RetroAnalogAxis, value: int):
        """Set analog stick axis value
        
        Args:
            port: Controller port (0-3)
            stick: Which stick (0=left, 1=right) or RetroAnalogStick enum
            axis: Which axis (0=X, 1=Y) or RetroAnalogAxis enum
            value: Analog value (-32768 to 32767, 0 is center)
        
        Example:
            emu.set_analog(0, RetroAnalogStick.LEFT, RetroAnalogAxis.X, 32767)
            emu.set_analog(0, 0, 0, 32767)  # Also valid
        """
        _ra_wrapper.set_analog(port, int(stick), int(axis), value)

    def clear_input(self):
        """Clear all input state (release all buttons, center all sticks)"""
        _ra_wrapper.clear_input()

    def _frame_generator(self):
        while True:
            self.step()
            yield self.get_video_frame(), self.get_audio_frame()
    
    def _video_frame_generator(self):
        while True:
            self.step()
            yield self.get_video_frame()

    def _audio_frame_generator(self):
        while True:
            self.step()
            yield self.get_audio_frame()
    
    @property
    def frames(self):
        if not self._game_loaded:
            raise RuntimeError("No game loaded")
        return self._frame_generator()

    @property
    def video_frames(self):
        if not self._game_loaded:
            raise RuntimeError("No game loaded")
        return self._video_frame_generator()

    @property
    def audio_frames(self):
        if not self._game_loaded:
            raise RuntimeError("No game loaded")
        return self._audio_frame_generator()

    def __del__(self):
        try:
            _ra_wrapper.shutdown()
        except:
            pass