import pygame
import math
import random
import os
import sys
from dataclasses import dataclass
from typing import List, Tuple, Optional

# --- FIX: FORCE GAME TO LOOK IN THE SCRIPT'S FOLDER ---
if getattr(sys, 'frozen', False):
    base_path = sys._MEIPASS
else:
    base_path = os.path.dirname(os.path.abspath(__file__))
os.chdir(base_path)

# --- Configuration ---
SCREEN_WIDTH = 1280 
SCREEN_HEIGHT = 720 
FPS = 60

# Room Settings
ROOM_COUNT = 8
WORLD_WIDTH = 14000
ROOM_WIDTH = WORLD_WIDTH // ROOM_COUNT
FLOOR_Y = SCREEN_HEIGHT - 100

# Timers (in Frames)
TIME_LIMIT_AIR = 3.5 * FPS        
TIME_LIMIT_POSSESSION = 5.5 * FPS

# Colors
BG_DARK = (15, 15, 25)
BG_ERROR = (50, 0, 0)
FLOOR_COLOR = (30, 30, 45)
WALL_COLOR = (60, 60, 80)
ROOM_LINE_COLOR = (50, 50, 70)
MASK_COLOR = (255, 215, 0)
MASK_INACTIVE = (150, 150, 150)
NPC_NORMAL = (100, 180, 255)
NPC_TARGET = (100, 255, 150)
NPC_USED = (80, 80, 80)      
NPC_PANIC = (255, 150, 255)
OBSTACLE_COLOR = (200, 50, 50)
BAR_AIR_COLOR = (0, 255, 255)
BAR_POSSESS_COLOR = (255, 215, 0)
TEXT_COLOR = (200, 200, 200)
TITLE_COLOR = (50, 255, 50)
ERROR_COLOR = (255, 0, 0)

# --- Dimensions ---
MASK_RADIUS = 28            
PLAYER_SIZE = MASK_RADIUS * 2
NPC_WIDTH = 110             
NPC_HEIGHT = 200            
HEAD_HITBOX_HEIGHT = 70      

# Physics
GRAVITY = 0.9               
TERMINAL_VELOCITY = 20      
MASK_FRICTION_AIR = 0.99
MASK_FRICTION_GROUND = 0.7
NPC_SPEED = 1                
NPC_PANIC_SPEED = 4          
LAUNCH_MAX_POWER = 25       

# --- MANUAL LEVEL DESIGN ---
LEVEL_LAYOUT = [
    [(0.5, 2, False)],
    [(0.3, 1, False), (0.8, 2, False)],
    [(0.4, 1, False)],
    [(0.3, 2, False), (0.7, 1, True)],
    [(0.2, 1, False), (0.5, 2, False), (0.8, 1, False)],
    [(0.3, 2, True), (0.7, 1, False)],
    [(0.2, 1, True), (0.5, 1, True), (0.8, 2, False)],
    [(0.2, 2, True), (0.5, 1, False), (0.8, 2, True)],
]

@dataclass
class Vector2:
    x: float
    y: float
   
    def __add__(self, other):
        return Vector2(self.x + other.x, self.y + other.y)
   
    def magnitude(self):
        return math.sqrt(self.x**2 + self.y**2)

class SoundManager:
    def __init__(self):
        self.sounds = {}
        self.enabled = True
        self.music_loaded = False
        
        try:
            # Force standard audio settings
            pygame.mixer.pre_init(44100, -16, 2, 512)
            pygame.mixer.init()
            
            # --- SOUND EFFECTS (Only the ones you wanted) ---
            files = {
                'jump': 'jump.mp3',
                'win': 'win.mp3',
                'game_over': 'game_over.mp3'
            }

            for name, filename in files.items():
                try:
                    if os.path.exists(filename):
                        self.sounds[name] = pygame.mixer.Sound(filename)
                        self.sounds[name].set_volume(1.0) # Max Volume
                        print(f"Loaded: {filename}")
                    else:
                        print(f"Warning: {filename} missing.")
                except Exception as e:
                    print(f"Error loading {filename}: {e}")

            # --- BACKGROUND MUSIC ---
            try:
                if os.path.exists('game_soundtrack.mp3'):
                    pygame.mixer.music.load('game_soundtrack.mp3')
                    pygame.mixer.music.set_volume(1.0) 
                    self.music_loaded = True
            except Exception as e:
                print(f"Music Error: {e}")

        except Exception as e:
            print(f"Audio Error: {e}")
            self.enabled = False

    def play(self, name):
        if self.enabled and name in self.sounds:
            self.sounds[name].play()

    def play_music(self):
        if self.enabled and self.music_loaded:
            try:
                if not pygame.mixer.music.get_busy():
                    pygame.mixer.music.play(-1)
            except: pass

    def stop_music(self):
        if self.enabled:
            pygame.mixer.music.stop()

class Obstacle:
    def __init__(self, x: float, y: float, w: float, h: float, is_deadly: bool = False, obs_type: int = 1, use_image: bool = True):
        self.rect = pygame.Rect(x, y, w, h)
        self.is_deadly = is_deadly
        self.color = OBSTACLE_COLOR if is_deadly else WALL_COLOR
        self.image = None
        self.has_image = False

        if use_image:
            filename = ""
            try:
                if obs_type == 1: 
                    filename = "computer_bad.png" if is_deadly else "computer.png"
                else: 
                    filename = "desk_bad.png" if is_deadly else "desk.png"
                
                if os.path.exists(filename):
                    raw_img = pygame.image.load(filename).convert_alpha()
                    self.image = pygame.transform.scale(raw_img, (w, h))
                    self.has_image = True
            except:
                self.has_image = False
   
    def draw(self, surface: pygame.Surface, camera_x: float):
        screen_x = self.rect.x - camera_x
        if -self.rect.width < screen_x < SCREEN_WIDTH:
            if self.has_image:
                surface.blit(self.image, (screen_x, self.rect.y))
            else:
                draw_rect = pygame.Rect(screen_x, self.rect.y, self.rect.width, self.rect.height)
                pygame.draw.rect(surface, self.color, draw_rect)
                pygame.draw.rect(surface, (0, 0, 0), draw_rect, 2)

class NPC:
    def __init__(self, x: float, is_target: bool = False):
        self.rect = pygame.Rect(x, FLOOR_Y - NPC_HEIGHT, NPC_WIDTH, NPC_HEIGHT)
        self.is_target = is_target
        self.color = NPC_TARGET if is_target else NPC_NORMAL
        self.wander_dir = 1 
        self.idle_timer = 0
        self.panic_timer = 0
        self.was_used = False
        self.frozen = False       
        self.sensor_timer = random.randint(0, 30) 

        self.frames = []
        self.frame_index = 0
        self.animation_speed_possessed = 0.15 
        self.animation_speed_ai = 0.05        
        self.has_images = False
        
        try:
            for i in range(1, 5):
                filename = f"char{i}.png"
                if os.path.exists(filename):
                    raw_img = pygame.image.load(filename).convert_alpha()
                    try:
                        scaled_img = pygame.transform.smoothscale(raw_img, (NPC_WIDTH, NPC_HEIGHT))
                    except:
                        scaled_img = pygame.transform.scale(raw_img, (NPC_WIDTH, NPC_HEIGHT))
                    self.frames.append(scaled_img)
            
            if len(self.frames) > 0:
                self.has_images = True
        except:
            self.has_images = False

    def trigger_panic(self, danger_x: float):
        if self.was_used or self.frozen: return
        self.panic_timer = 120 
        if self.rect.centerx > danger_x:
            self.wander_dir = 1
        else:
            self.wander_dir = -1

    def update(self, obstacles: List[Obstacle], player_x: float, all_npcs: List['NPC']) -> bool:
        if self.was_used: return False

        self.sensor_timer += 1
        if self.sensor_timer > 30:
            self.sensor_timer = 0
            self.frozen = self.check_space_too_small(obstacles)
            if self.check_should_cull(player_x, all_npcs): return True 

        if self.frozen:
            if self.panic_timer > 0: self.panic_timer -= 1
            return False

        speed = NPC_SPEED
        if self.panic_timer > 0:
            self.panic_timer -= 1
            speed = NPC_PANIC_SPEED 
        else:
            if self.idle_timer > 0:
                self.idle_timer -= 1
                return False

        move_x = self.wander_dir * speed
        self.rect.x += move_x
        
        for obs in obstacles:
            if self.rect.colliderect(obs.rect):
                self.rect.x -= move_x 
                self.wander_dir *= -1 
                if self.panic_timer == 0:
                    self.idle_timer = 60 
                break
        return False 

    def check_space_too_small(self, obstacles):
        relevant_obs = [o.rect for o in obstacles if o.rect.top < self.rect.bottom and o.rect.bottom > self.rect.top]
        left_walls = [r.right for r in relevant_obs if r.right < self.rect.left]
        dist_left = (self.rect.left - max(left_walls)) if left_walls else 1000
        right_walls = [r.left for r in relevant_obs if r.left > self.rect.right]
        dist_right = (min(right_walls) - self.rect.right) if right_walls else 1000
        total_room_width = dist_left + dist_right + self.rect.width
        return total_room_width < (NPC_WIDTH * 1.5)

    def check_should_cull(self, player_x, all_npcs):
        if self.is_target: return False 
        dist_to_player = abs(self.rect.centerx - player_x)
        if dist_to_player > 2000:
            neighbors = 0
            for other in all_npcs:
                if other is not self and abs(other.rect.x - self.rect.x) < 300:
                    neighbors += 1
            if neighbors > 2: return True
        return False

    def draw(self, surface: pygame.Surface, camera_x: float, possessed: bool = False):
        screen_x = self.rect.x - camera_x
        if -200 < screen_x < SCREEN_WIDTH + 200:
            is_moving = True 
            current_speed = self.animation_speed_ai

            if possessed:
                is_moving = (self.idle_timer == 0)
                current_speed = self.animation_speed_possessed
            else:
                is_moving = (not self.frozen and self.idle_timer == 0)

            if self.has_images:
                if is_moving:
                    self.frame_index += current_speed
                    if self.frame_index >= len(self.frames): self.frame_index = 0
                else:
                    self.frame_index = 0 
                
                idx = int(self.frame_index) % len(self.frames)
                current_img = self.frames[idx].copy()

                if possessed:
                    current_img.fill((255, 215, 0, 100), special_flags=pygame.BLEND_RGBA_MULT)
                elif self.was_used:
                    current_img.fill((80, 80, 80, 255), special_flags=pygame.BLEND_RGBA_MULT)
                elif self.panic_timer > 0:
                    current_img.fill((255, 100, 100, 255), special_flags=pygame.BLEND_RGBA_MULT)
                elif self.frozen:
                    current_img.fill((100, 100, 150, 255), special_flags=pygame.BLEND_RGBA_MULT)
                elif self.is_target:
                    current_img.fill((200, 255, 200, 255), special_flags=pygame.BLEND_RGBA_MULT)

                if self.wander_dir == -1:
                    current_img = pygame.transform.flip(current_img, True, False)

                surface.blit(current_img, (screen_x, self.rect.y))
            else:
                draw_rect = pygame.Rect(screen_x, self.rect.y, self.rect.width, self.rect.height)
                if possessed: body_color = (255, 215, 0) 
                elif self.was_used: body_color = NPC_USED 
                elif self.panic_timer > 0: body_color = NPC_PANIC   
                elif self.frozen: body_color = (100, 100, 150)
                else: body_color = self.color
                pygame.draw.rect(surface, body_color, draw_rect, border_radius=10)
                pygame.draw.rect(surface, (0, 0, 0), draw_rect, 3, border_radius=10)

            head_y = self.rect.y + HEAD_HITBOX_HEIGHT - 4
            pygame.draw.line(surface, (255,0,0), (screen_x, head_y), (screen_x + self.rect.width, head_y), 2)

class Mask:
    def __init__(self, x: float, y: float, sound_manager):
        self.pos = Vector2(x, y)
        self.vel = Vector2(0, 0)
        self.radius = MASK_RADIUS
        self.possessed_npc: Optional[NPC] = None
        self.last_possessed_npc: Optional[NPC] = None
        self.cooldown_timer = 0
        self.can_launch = True
        self.is_dragging = False
        self.air_timer = TIME_LIMIT_AIR
        self.possession_timer = TIME_LIMIT_POSSESSION
        self.sounds = sound_manager

    def get_rect(self) -> pygame.Rect:
        return pygame.Rect(self.pos.x - self.radius, self.pos.y - self.radius, self.radius * 2, self.radius * 2)

    def launch(self, mouse_world_x: float, mouse_y: float):
        dx = self.pos.x - mouse_world_x
        dy = self.pos.y - mouse_y
        dist = math.sqrt(dx**2 + dy**2)
        if dist > 5 and self.can_launch:
            # Play jump sound
            self.sounds.play('jump')
            
            if self.possessed_npc:
                self.possessed_npc.was_used = True
                self.last_possessed_npc = self.possessed_npc
                self.cooldown_timer = 40 
                self.air_timer = TIME_LIMIT_AIR
                self.pos.y -= 60 
            else:
                self.pos.y -= 10
            self.possessed_npc = None
            power = min(dist * 0.2, LAUNCH_MAX_POWER)
            angle = math.atan2(dy, dx)
            self.vel.x = math.cos(angle) * power
            self.vel.y = math.sin(angle) * power
            self.can_launch = False
            self.is_dragging = False

    def update(self, obstacles: List[Obstacle], npcs: List[NPC]):
        if self.cooldown_timer > 0: self.cooldown_timer -= 1
        else: self.last_possessed_npc = None

        if self.possessed_npc:
            self.pos.x = self.possessed_npc.rect.centerx
            self.pos.y = self.possessed_npc.rect.y + 10
            self.can_launch = True
            self.air_timer = TIME_LIMIT_AIR
            self.possession_timer -= 1
            if self.possession_timer <= 0: return True 
            return False 

        self.air_timer -= 1
        if self.air_timer <= 0: return True 

        self.vel.y = min(self.vel.y + GRAVITY, TERMINAL_VELOCITY)
        self.pos.x += self.vel.x
        self.pos.y += self.vel.y
        self.vel.x *= MASK_FRICTION_AIR
        self.vel.y *= MASK_FRICTION_AIR

        if self.pos.y > FLOOR_Y - self.radius:
            self.pos.y = FLOOR_Y - self.radius
            self.vel.y = 0 
            self.vel.x = 0  
            self.can_launch = True

        mask_rect = self.get_rect()
        for obs in obstacles:
            if mask_rect.colliderect(obs.rect):
                if obs.is_deadly: return True 
                
                overlap_left = mask_rect.right - obs.rect.left
                overlap_right = obs.rect.right - mask_rect.left
                overlap_top = mask_rect.bottom - obs.rect.top
                overlap_bottom = obs.rect.bottom - mask_rect.top
                min_overlap = min(overlap_left, overlap_right, overlap_top, overlap_bottom)

                if min_overlap == overlap_left:
                    self.pos.x = obs.rect.left - self.radius
                    self.vel.x = 0
                elif min_overlap == overlap_right:
                    self.pos.x = obs.rect.right + self.radius
                    self.vel.x = 0
                elif min_overlap == overlap_top:
                    self.pos.y = obs.rect.top - self.radius
                    self.vel.y = 0
                    self.vel.x = 0 
                    self.can_launch = True
                elif min_overlap == overlap_bottom:
                    self.pos.y = obs.rect.bottom + self.radius
                    self.vel.y = 0
                mask_rect = self.get_rect()

        if not self.possessed_npc and self.vel.magnitude() > 2:
            for npc in npcs:
                if not npc.is_target and mask_rect.colliderect(npc.rect):
                    if npc.was_used or npc == self.last_possessed_npc:
                        self.vel.x *= -0.5
                        self.vel.y *= -0.5
                        self.pos.y -= 5
                        continue 
                    
                    hit_y_relative = self.pos.y - npc.rect.y
                    if hit_y_relative < (HEAD_HITBOX_HEIGHT - 4):
                        self.possessed_npc = npc
                        # Removed possess sound as requested
                        self.vel = Vector2(0, 0)
                        self.possession_timer = TIME_LIMIT_POSSESSION 
                        self.pos.x = npc.rect.centerx
                        self.pos.y = npc.rect.y + 10
                        break
                    else:
                        npc.trigger_panic(self.pos.x)
                        self.vel.x *= -0.5
                        self.vel.y *= -0.5
                        self.pos.y -= 5 
                        self.last_possessed_npc = npc 
                        self.cooldown_timer = 20
                        break
        return False

    def draw(self, surface: pygame.Surface, camera_x: float):
        screen_pos = (int(self.pos.x - camera_x), int(self.pos.y))
        
        try:
            mask_img = pygame.image.load("mask_tentacles.png").convert_alpha()
            mask_img = pygame.transform.scale(mask_img, (self.radius*2, self.radius*2))
            surface.blit(mask_img, (screen_pos[0]-self.radius, screen_pos[1]-self.radius))
        except:
            # FALLBACK MASK
            offset_x = 0
            if self.vel.magnitude() > 1:
                trail_angle = math.atan2(self.vel.y, self.vel.x) + math.pi 
                offset_x = math.cos(trail_angle) * 20
            
            for i in range(3):
                start_x = screen_pos[0] + (i - 1) * 8
                start_y = screen_pos[1]
                wiggle = math.sin(pygame.time.get_ticks() * 0.01 + i) * 5
                end_x = start_x + offset_x + wiggle
                end_y = start_y + 30 
                pygame.draw.line(surface, (255, 100, 100), (start_x, start_y), (end_x, end_y), 3)

            mask_rect = pygame.Rect(screen_pos[0]-20, screen_pos[1]-15, 40, 30)
            pygame.draw.ellipse(surface, (230, 230, 220), mask_rect) 
            pygame.draw.ellipse(surface, (0, 0, 0), mask_rect, 2) 

        bar_w = 60
        bar_h = 8
        bar_x = screen_pos[0] - bar_w // 2
        bar_y = screen_pos[1] - self.radius - 20 
        
        if self.possessed_npc:
            ratio = max(0, self.possession_timer / TIME_LIMIT_POSSESSION)
            pygame.draw.rect(surface, (50, 50, 50), (bar_x, bar_y, bar_w, bar_h))
            pygame.draw.rect(surface, BAR_POSSESS_COLOR, (bar_x, bar_y, bar_w * ratio, bar_h))
        else:
            ratio = max(0, self.air_timer / TIME_LIMIT_AIR)
            pygame.draw.rect(surface, (50, 50, 50), (bar_x, bar_y, bar_w, bar_h))
            pygame.draw.rect(surface, BAR_AIR_COLOR, (bar_x, bar_y, bar_w * ratio, bar_h))

def generate_level():
    obstacles = []
    npcs = []
    obstacles.append(Obstacle(0, 0, 50, SCREEN_HEIGHT, use_image=False))
    obstacles.append(Obstacle(WORLD_WIDTH - 50, 0, 50, SCREEN_HEIGHT, use_image=False))
    for room_idx, room_data in enumerate(LEVEL_LAYOUT):
        room_x_start = room_idx * ROOM_WIDTH
        for (pos_pct, obs_type, is_deadly) in room_data:
            if obs_type == 1: w, h = 120, 280
            else: w, h = 224, 112
            x = room_x_start + int(pos_pct * ROOM_WIDTH)
            obstacles.append(Obstacle(x, FLOOR_Y - h, w, h, is_deadly, obs_type))
    obstacles.sort(key=lambda ob: ob.rect.x)
    previous_right_x = 50 
    npc_count = 0 
    for obs in obstacles:
        current_left_x = obs.rect.x
        gap_size = current_left_x - previous_right_x
        if gap_size > (PLAYER_SIZE * 2.5):
            npc_count += 1
            if npc_count % 2 == 0:
                spawn_x = current_left_x - NPC_WIDTH - PLAYER_SIZE
                npc = NPC(spawn_x)
                npc.wander_dir = -1 
            else:
                spawn_x = previous_right_x + (gap_size // 2) - (NPC_WIDTH // 2)
                npc = NPC(spawn_x)
                npc.wander_dir = -1 
            npcs.append(npc)
        previous_right_x = obs.rect.right
    remaining_gap = (WORLD_WIDTH - 400) - previous_right_x
    if remaining_gap > (PLAYER_SIZE * 3):
        spawn_x = previous_right_x + (remaining_gap // 2) - (NPC_WIDTH // 2)
        npcs.append(NPC(spawn_x))
    target = NPC(WORLD_WIDTH - 400, is_target=True)
    npcs.append(target)
    return obstacles, npcs, target

class Game:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("Subject 0: The goo")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font(None, 60)
        self.small_font = pygame.font.Font(None, 36)
        self.win_start_time = 0 
        self.has_played_win_sound = False
        
        self.sounds = SoundManager()

        try:
            bg_raw = pygame.image.load("mask science lab wall.png").convert()
            scale_ratio = SCREEN_HEIGHT / bg_raw.get_height()
            new_bg_w = int(bg_raw.get_width() * scale_ratio)
            self.bg_img = pygame.transform.scale(bg_raw, (new_bg_w, SCREEN_HEIGHT))
            
            floor_raw = pygame.image.load("mask science lab floor.png").convert()
            floor_h = SCREEN_HEIGHT - FLOOR_Y
            scale_ratio_floor = floor_h / floor_raw.get_height()
            new_floor_w = int(floor_raw.get_width() * scale_ratio_floor)
            self.floor_img = pygame.transform.scale(floor_raw, (new_floor_w, floor_h))
            self.has_graphics = True
        except:
            self.has_graphics = False

        self.reset()
        self.state = "story" 

    def reset(self):
        self.obstacles, self.npcs, self.target = generate_level()
        self.mask = Mask(200, FLOOR_Y - 50, self.sounds)
        self.camera_x = 0
        self.has_played_win_sound = False
        if hasattr(self, 'state') and self.state != "story":
            self.state = "playing"
            self.sounds.play_music() 

    def draw_text_centered(self, text, font, color, y_offset):
        text_surf = font.render(text, True, color)
        rect = text_surf.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + y_offset))
        self.screen.blit(text_surf, rect)

    def run(self):
        while True:
            self.clock.tick(FPS)
            mouse_pos = pygame.mouse.get_pos()
            world_mouse_x = mouse_pos[0] + self.camera_x

            for event in pygame.event.get():
                if event.type == pygame.QUIT: return
                
                if self.state == "story":
                    if event.type == pygame.KEYDOWN and event.key == pygame.K_SPACE:
                        self.state = "playing"
                        self.reset()
                        self.sounds.play_music() 
                
                elif self.state in ["dead", "won"]:
                    if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                        self.reset()
                        self.state = "playing"

                elif self.state == "playing":
                    if event.type == pygame.KEYDOWN and event.key == pygame.K_r: 
                        self.reset()
                    
                    if event.type == pygame.MOUSEBUTTONDOWN:
                        if self.mask.can_launch: self.mask.is_dragging = True
                    if event.type == pygame.MOUSEBUTTONUP and self.mask.is_dragging:
                        self.mask.is_dragging = False
                        self.mask.launch(world_mouse_x, mouse_pos[1])

            self.screen.fill(BG_DARK)

            if self.state == "story":
                self.draw_text_centered("SUBJECT 0: NEURO-PARASITE", self.font, TITLE_COLOR, -150)
                
                story_lines = [
                    "Experiment 734 failed.",
                    "Your consciousness has transferred to the prototype Neural Parasite Mask.",
                    "Your original body lies comatose in Sector Z.",
                    "Your tentacles wither without a host's neural energy.",
                    "Latch onto scientists to hijack their motor functions.",
                    "The host body will eventually reject you.",
                    "Reach your body before you succumb.",
                    "",
                    "[ PRESS SPACE TO BEGIN ]"
                ]
                
                for i, line in enumerate(story_lines):
                    self.draw_text_centered(line, self.small_font, TEXT_COLOR, -60 + (i * 40))
                
                cx, cy = SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 200
                pygame.draw.ellipse(self.screen, (230, 230, 220), (cx-40, cy-30, 80, 60))
                for i in range(3):
                    pygame.draw.line(self.screen, (255, 100, 100), (cx + (i-1)*20, cy+30), (cx + (i-1)*25, cy+70), 5)

            elif self.state == "playing":
                target_cam = self.mask.pos.x - SCREEN_WIDTH // 3
                self.camera_x += (target_cam - self.camera_x) * 0.1
                self.camera_x = max(0, min(self.camera_x, WORLD_WIDTH - SCREEN_WIDTH))

                if self.mask.update(self.obstacles, self.npcs):
                    self.state = "dead"
                    self.sounds.stop_music()
                    self.sounds.play('game_over')
                
                if self.mask.possessed_npc:
                    keys = pygame.key.get_pressed()
                    dx = (keys[pygame.K_d] - keys[pygame.K_a]) * 4 
                    
                    if dx > 0:
                        self.mask.possessed_npc.wander_dir = 1
                        self.mask.possessed_npc.idle_timer = 0
                    elif dx < 0:
                        self.mask.possessed_npc.wander_dir = -1
                        self.mask.possessed_npc.idle_timer = 0
                    else:
                        self.mask.possessed_npc.idle_timer = 2
                    
                    self.mask.possessed_npc.rect.x += dx
                    for obs in self.obstacles:
                        if self.mask.possessed_npc.rect.colliderect(obs.rect):
                            if obs.is_deadly: 
                                self.state = "dead"
                                self.sounds.stop_music()
                                self.sounds.play('game_over') 
                            self.mask.possessed_npc.rect.x -= dx

                for npc in self.npcs[:]:
                    if npc != self.mask.possessed_npc:
                        should_delete = npc.update(self.obstacles, self.mask.pos.x, self.npcs)
                        if should_delete: self.npcs.remove(npc)

                if self.mask.get_rect().colliderect(self.target.rect):
                    self.state = "won"
                    self.win_start_time = pygame.time.get_ticks() 
                    self.sounds.stop_music()
                    if not self.has_played_win_sound:
                        self.sounds.play('win')
                        self.has_played_win_sound = True

                if self.has_graphics:
                    bg_w = self.bg_img.get_width()
                    start_x = -(self.camera_x % bg_w)
                    tiles_needed = math.ceil(SCREEN_WIDTH / bg_w) + 1
                    for i in range(tiles_needed):
                        self.screen.blit(self.bg_img, (start_x + (i * bg_w), 0))
                else:
                    self.screen.fill(BG_DARK)

                for i in range(1, ROOM_COUNT):
                    line_x = (i * ROOM_WIDTH) - self.camera_x
                    if 0 < line_x < SCREEN_WIDTH:
                        pygame.draw.line(self.screen, ROOM_LINE_COLOR, (line_x, 0), (line_x, SCREEN_HEIGHT), 1)

                if self.has_graphics:
                    floor_w = self.floor_img.get_width()
                    start_x = -(self.camera_x % floor_w)
                    tiles_needed = math.ceil(SCREEN_WIDTH / floor_w) + 1
                    for i in range(tiles_needed):
                        self.screen.blit(self.floor_img, (start_x + (i * floor_w), FLOOR_Y))
                else:
                    pygame.draw.rect(self.screen, FLOOR_COLOR, (0 - self.camera_x, FLOOR_Y, WORLD_WIDTH, 100))
                
                for obs in self.obstacles: obs.draw(self.screen, self.camera_x)
                for npc in self.npcs: npc.draw(self.screen, self.camera_x, npc == self.mask.possessed_npc)
                self.mask.draw(self.screen, self.camera_x)
                
                if self.mask.is_dragging:
                    pygame.draw.line(self.screen, (255,255,255), (self.mask.pos.x - self.camera_x, self.mask.pos.y), mouse_pos, 4)

            elif self.state == "dead":
                if self.mask.possessed_npc:
                    msg = "HOST REJECTION CRITICAL"
                    sub = "Biomass Destroyed. Press [R] to Restart."
                else:
                    msg = "NO HOST FOUND"
                    sub = "Biomass Dried Out. Press [R] to Restart."

                self.draw_text_centered(msg, self.font, (255, 50, 50), -50)
                self.draw_text_centered(sub, self.small_font, TEXT_COLOR, 50)

            elif self.state == "won":
                # --- TWIST LOGIC ---
                time_passed = pygame.time.get_ticks() - self.win_start_time

                if time_passed < 3000:
                    self.draw_text_centered("SUBJECT 0 LOCATED", self.font, TITLE_COLOR, -50)
                    self.draw_text_centered("INITIATING NEURAL TRANSFER...", self.small_font, TEXT_COLOR, 10)
                else:
                    self.screen.fill(BG_ERROR) 
                    self.draw_text_centered("CRITICAL ERROR", self.font, (0, 0, 0), -50)
                    self.draw_text_centered("ORIGINAL HOST INCOMPATIBLE.", self.small_font, (0, 0, 0), 10)
                    self.draw_text_centered("TRANSFER REJECTED.", self.small_font, (0, 0, 0), 50)
                    self.draw_text_centered("[R] TO RESTART SIMULATION", self.small_font, (255, 255, 255), 150)

            pygame.display.flip()

if __name__ == "__main__":
    game = Game()
    game.run()
    pygame.quit()