import gymnasium as gym
from gymnasium import spaces
import numpy as np
from vision_agent import EldenVision
from input_driver import EldenController
import time
from evdev import ecodes as e

class EldenRingEnv(gym.Env):
    def __init__(self):
        super(EldenRingEnv, self).__init__()
        self.vision = EldenVision(stack_size=4)
        self.controller = EldenController()
        self.observation_space = spaces.Box(low=0, high=255, shape=(640, 640, 12), dtype=np.uint8)
        self.action_space = spaces.MultiDiscrete([4, 5, 3, 3, 2])
        
        self.last_p_hp = 100.0
        self.last_b_hp = 100.0
        self.frame_skip = 4 

    def step(self, actions):
        combat, move, cam_x, cam_y, switch = actions
        if switch == 1: self.controller.switch_item()
        if combat == 1: self.controller.dodge()
        elif combat == 2: self.controller.attack()
        elif combat == 3: self.controller.use_item()
        self.controller.move(move)
        
        mx = -50 if cam_x == 1 else (50 if cam_x == 2 else 0)
        my = -50 if cam_y == 1 else (50 if cam_y == 2 else 0)
        if mx != 0 or my != 0: self.controller.look(mx, my)

        for _ in range(self.frame_skip):
            obs, hp, fp, sp, boss = self.vision.get_state()

        reward = 0.01
        if combat == 3 and hp > 90: reward -= 5.0
        if hp < self.last_p_hp: reward -= 20.0
        if boss < self.last_b_hp: reward += 10.0 * ((self.last_b_hp - boss) / 5.0)

        self.last_p_hp, self.last_b_hp = hp, boss
        done = hp <= 0 or (boss <= 0 and self.last_b_hp > 10)
        return obs, reward, done, False, {}

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        print(f"\n[NAV] DEATH DETECTED. Waiting 25s for Respawn...")
        time.sleep(25) 
        
        print("[NAV] Running to Maliketh Fog Gate...")
        
        # 1. Exit Grace Room
        self.controller.move(1) # W
        time.sleep(2.0)
        
        # 2. Turn Right to stairs
        self.controller.look(1800, 0) 
        time.sleep(0.5)
        
        # 3. Run up the stairs
        self.controller.move(1) 
        time.sleep(4.5)
        
        # 4. Turn Left to bridge
        self.controller.look(-1600, 0)
        time.sleep(0.5)
        
        # 5. Sprint across bridge and look for fog
        print("[NAV] Sprinting across bridge...")
        start_nav = time.time()
        found_fog = False
        
        while time.time() - start_nav < 25: # 25s timeout for navigation
            with self.vision.lock:
                hsv = self.vision.frame_hsv.copy()
            
            if self.vision.detect_fog(hsv):
                print("[NAV] FOG DETECTED! Entering Arena...")
                self.controller.move(0) 
                time.sleep(0.2)
                self.controller.press(e.KEY_E) # Interact
                time.sleep(5.0) # Extra safety for entry animation
                found_fog = True
                break
                
            self.controller.move(1) # Keep running forward
            time.sleep(0.1)

        if not found_fog:
            print("[NAV] Warning: Fog not found. Sentinel might have blocked us.")
        
        self.controller.move(0)
        obs, hp, fp, sp, boss = self.vision.get_state()
        self.last_p_hp, self.last_b_hp = hp, boss
        return obs, {}