import math
import random
import pyxel
from enum import Enum, auto

SCREEN_W = 128
SCREEN_H = 144
PLAYER_SPEED = 1.2
GRAVITY = 0.25
JUMP_SPEED = -4.5
PLATFORM_H = 5


# ============================================================
#  REWARD TIMER
# ============================================================

class RewardTimer:
    def __init__(self):
        self.time_left = 0
        self.active = False
        self.reward = 0

    def start(self, minutes, reward):
        self.time_left = minutes * 60
        self.reward = reward
        self.active = True

    def update(self, dt):
        if self.active:
            self.time_left -= dt
            if self.time_left <= 0:
                self.time_left = 0
                self.active = False
                return self.reward
        return 0

    def draw(self):
        if self.active:
            mins = int(self.time_left // 60)
            secs = int(self.time_left % 60)
            label = f"TIMER {mins:02d}:{secs:02d}"
            w = len(label) * 4 + 4
            pyxel.rect(2, 31, w + 2, 11, 1)
            pyxel.rectb(2, 31, w + 2, 11, 7)
            pyxel.text(4, 33, label, 10)


# ============================================================
#  PLAYER
# ============================================================

class Player:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.dir = "down"
        self.moving = False

    def update(self, dt, dx, dy):
        self.moving = dx != 0 or dy != 0

        if dx > 0:
            self.dir = "right"
        elif dx < 0:
            self.dir = "left"
        elif dy != 0:
            self.dir = "down"

        self.x += dx * PLAYER_SPEED * 60 * dt
        self.y += dy * PLAYER_SPEED * 60 * dt

        self.x = max(8, min(SCREEN_W - 24, self.x))
        self.y = max(48, min(SCREEN_H - 24, self.y))

    def draw(self):
        frame = (pyxel.frame_count // 8) % 2
        u, v, w, h = 0, 0, 16, 16

        if self.dir == "down":
            u = (16 if frame == 0 else 32) if self.moving else 0

        elif self.dir == "right":
            u = (64 if frame == 0 else 80) if self.moving else 48
            w = 16

        elif self.dir == "left":
            u = (64 if frame == 0 else 80) if self.moving else 48
            w = -16  # flip horizontally

        pyxel.blt(int(self.x), int(self.y), 0, u, v, w, h, 2)


# ============================================================
#  CAT (unchanged)
# ============================================================

class CatState(Enum):
    IDLE = auto()
    WALK = auto()
    FOLLOW = auto()
    SLEEP = auto()


class Cat:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.state = CatState.SLEEP
        self.timer = random.uniform(5, 12)
        self.walk_target = None

    def update(self, dt, player):
        self.timer -= dt

        if self.state == CatState.SLEEP:
            if self.timer <= 0:
                self.state = CatState.IDLE
                self.timer = random.uniform(6, 12)
            return

        if self.state == CatState.IDLE:
            if self.timer <= 0:
                if random.random() < 0.4:
                    self.state = CatState.WALK
                    self.timer = random.uniform(4, 8)
                    self.walk_target = (
                        player.x + random.randint(-40, 40),
                        player.y + random.randint(-10, 10)
                    )
                else:
                    self.state = CatState.SLEEP
                    self.timer = random.uniform(8, 16)
            return

        if self.state == CatState.WALK:
            tx, ty = self.walk_target
            dx = tx - self.x
            dy = ty - self.y
            dist = math.hypot(dx, dy)

            if dist > 1:
                self.x += dx / dist * 30 * dt
                self.y += dy / dist * 30 * dt

            if dist <= 1 or self.timer <= 0:
                self.state = CatState.IDLE
                self.timer = random.uniform(6, 12)
            return

        if self.state == CatState.FOLLOW:
            dx = player.x - self.x
            dy = player.y - self.y
            dist = math.hypot(dx, dy)

            if dist > 10:
                self.x += dx / dist * 40 * dt
                self.y += dy / dist * 40 * dt

            return

    def start_follow(self):
        self.state = CatState.FOLLOW
        self.timer = 99999

    def is_near(self, px, py):
        return (self.x - px) ** 2 + (self.y - py) ** 2 < 20 ** 2

    def draw(self):
        u, v, w, h = 0, 32, 16, 16
        if self.state == CatState.SLEEP:
            u = 80
        elif self.state == CatState.FOLLOW:
            u = 48
        elif self.state == CatState.IDLE:
            blink_phase = (pyxel.frame_count // 6) % 80
            u = 64 if blink_phase < 4 else 0
        elif self.state == CatState.WALK:
            frame = (pyxel.frame_count // 10) % 2
            u = 16 if frame == 0 else 32
        pyxel.blt(int(self.x), int(self.y), 0, u, v, w, h, 2)


# ============================================================
#  DREAM LEVELS (3 LEVELS + THEMES)
# ============================================================

class DreamLevel:
    def __init__(self, level, app):
        self.level = level
        self.app = app
        self.load_level()
        self.reset()

    def load_level(self):
        if self.level == 0:
            self.platforms = [
                (0,126,28),(34,116,18),(56,106,16),(76,96,16),
                (50,86,16),(22,76,18),(44,66,16),(66,56,16),(88,46,40)
            ]
            self.lava_y = 132

        elif self.level == 1:
            self.platforms = [
                (0,126,20),(30,110,14),(60,95,14),(90,80,14),
                (70,65,14),(40,50,14),(10,35,14),(50,20,40)
            ]
            self.lava_y = 138

        elif self.level == 2:
            self.platforms = [
                (0,126,18),(40,115,12),(80,105,12),
                (20,90,10),(60,75,10),(100,60,10),
                (30,45,8),(70,30,8),(100,15,30)
            ]
            self.lava_y = 140

    def reset(self):
        self.px = float(self.platforms[0][0] + 4)
        self.py = float(self.platforms[0][1] - 16)
        self.vy = 0.0
        self.on_ground = False
        self.facing_right = True
        self.won = False
        self.won_timer = 0.0
        self.deaths = 0

    def update(self, dt):
        if self.won:
            self.won_timer -= dt
            return self.won_timer <= 0

        moving_right = pyxel.btn(pyxel.KEY_RIGHT) or pyxel.btn(pyxel.KEY_D)
        moving_left  = pyxel.btn(pyxel.KEY_LEFT)  or pyxel.btn(pyxel.KEY_A)
        jump = pyxel.btnp(pyxel.KEY_Z) or pyxel.btnp(pyxel.KEY_UP)

        dx = int(moving_right) - int(moving_left)
        if dx > 0:
            self.facing_right = True
        elif dx < 0:
            self.facing_right = False

        self.px += dx * 1.3
        self.px = max(0, min(SCREEN_W - 16, self.px))

        if jump and self.on_ground:
            self.vy = JUMP_SPEED

        self.vy = min(self.vy + GRAVITY, 6)
        self.py += self.vy

        self.on_ground = False
        if self.vy >= 0:
            for i, (plx, ply, plw) in enumerate(self.platforms):
                feet = self.py + 16
                if (plx <= self.px + 12 and self.px + 4 <= plx + plw and
                        ply <= feet <= ply + PLATFORM_H + self.vy + 2):
                    self.py = ply - 16
                    self.vy = 0
                    self.on_ground = True
                    if i == len(self.platforms) - 1:
                        self.won = True
                        self.won_timer = 2.5
                    break

        if self.py > self.lava_y:
            self.deaths += 1
            self.reset()

        return False

    def draw(self):
        theme = self.app.setting_theme

        if theme == 0:  # Classic
            bg = 2
            star1, star2 = 7, 6
            cloud = 13
            lava1, lava2 = 8, 9
        elif theme == 1:  # Sunset
            bg = 9
            star1, star2 = 10, 7
            cloud = 8
            lava1, lava2 = 2, 3
        else:  # Ocean
            bg = 1
            star1, star2 = 12, 11
            cloud = 5
            lava1, lava2 = 6, 7

        pyxel.cls(bg)

        # Stars
        for sx, sy in [
            (random.randint(0, SCREEN_W), random.randint(0, 90))
            for _ in range(35)
        ]:
            col = star1 if (pyxel.frame_count // 20 + sx) % 6 != 0 else star2
            pyxel.pset(sx, sy, col)

        # Clouds
        for i, cx in enumerate([10, 50, 90]):
            cy = 20 + int(math.sin((pyxel.frame_count / 40) + i * 2) * 3)
            pyxel.rect(cx, cy, 18, 6, cloud)
            pyxel.rect(cx + 3, cy - 3, 10, 5, cloud)

        # HUD
        pyxel.text(2, 2, "DREAMLAND", 7)
        if self.deaths > 0:
            pyxel.text(80, 2, f"Falls:{self.deaths}", 8)

        # Goal label
        gx, gy = self.platforms[-1][0], self.platforms[-1][1]
        pyxel.text(gx + 4, gy - 9, "GOAL!", 10)

        # Platforms
        for i, (plx, ply, plw) in enumerate(self.platforms):
            top_col = 11 if i == len(self.platforms) - 1 else 7
            fill_col = 3 if i < len(self.platforms) - 1 else 11
            pyxel.rect(plx, ply, plw, PLATFORM_H, fill_col)
            pyxel.rect(plx, ply, plw, 1, top_col)

        # Lava
        lava_col = lava1 if (pyxel.frame_count // 6) % 2 == 0 else lava2
        pyxel.rect(0, self.lava_y, SCREEN_W, SCREEN_H - self.lava_y, lava_col)

        # Player sprite
        moving = (pyxel.btn(pyxel.KEY_RIGHT) or pyxel.btn(pyxel.KEY_D) or
                  pyxel.btn(pyxel.KEY_LEFT)  or pyxel.btn(pyxel.KEY_A))
        frame = (pyxel.frame_count // 8) % 2
        su = (64 if frame == 0 else 80) if moving else 48
        sw = 16 if self.facing_right else -16
        pyxel.blt(int(self.px), int(self.py), 0, su, 0, sw, 16, 2)

        # Win overlay
        if self.won:
            pyxel.rect(14, 58, 100, 28, 1)
            pyxel.rectb(14, 58, 100, 28, 10)
            pyxel.text(22, 65, "YOU ESCAPED!", 10)
            pyxel.text(22, 75, "Waking up...", 7)


# ============================================================
#  MAIN APP (START)
# ============================================================

class App:
    def __init__(self):
        pyxel.init(SCREEN_W, SCREEN_H, title="Midnight Corner", fps=60)
        pyxel.load("res.pyxres")

        self.state = "title"
        self.menu_index = 0
        self.timer_index = 0
        self.settings_index = 0
        self.fade = 1.0

        self.setting_music = True
        self.setting_lighting = True
        self.setting_speed = 1
        self.setting_theme = 0  # NEW

        self.player = Player(60, 110)
        self.cat = Cat(80, 110)

        self.coins = 0
        self.message = ""
        self.message_timer = 0

        self.shop_open = False
        self.timer_menu_open = False

        self.reward_timer = RewardTimer()

        self.current_dream_level = 0
        self.dream = DreamLevel(self.current_dream_level, self)

        pyxel.playm(0, loop=True)
        pyxel.run(self.update, self.draw)
    # ---------------------------- UPDATE ----------------------------
    def update(self):
        dt = 1 / 60

        if self.state == "title":
            self.update_title(dt)
            return

        if self.state == "settings":
            self.update_settings(dt)
            return

        if self.state == "dream":
            done = self.dream.update(dt)
            if done:
                if self.dream.won:
                    self.coins += 5
                    self.message = "+5 coins! Sweet dreams!"
                    self.message_timer = 3

                    # Next dream level
                    self.current_dream_level = (self.current_dream_level + 1) % 3

                self.state = "game"
            return

        dx = int(pyxel.btn(pyxel.KEY_RIGHT) or pyxel.btn(pyxel.KEY_D)) - \
             int(pyxel.btn(pyxel.KEY_LEFT)  or pyxel.btn(pyxel.KEY_A))

        dy = int(pyxel.btn(pyxel.KEY_DOWN) or pyxel.btn(pyxel.KEY_S)) - \
             int(pyxel.btn(pyxel.KEY_UP)   or pyxel.btn(pyxel.KEY_W))

        speed_mult = [0.6, 1.0, 1.6][self.setting_speed]
        self.player.update(dt * speed_mult, dx, dy)
        self.cat.update(dt, self.player)

        reward = self.reward_timer.update(dt)
        if reward > 0:
            self.coins += reward
            self.message = f"+{reward} coins!"
            self.message_timer = 3

        if pyxel.btnp(pyxel.KEY_T) and not self.reward_timer.active:
            if self.shop_open:
                self.shop_open = False
            self.timer_menu_open = not self.timer_menu_open
            self.timer_index = 0

        if self.timer_menu_open:
            self.update_timer_menu()
            return

        if pyxel.btnp(pyxel.KEY_SPACE):
            if self.shop_open:
                self.shop_open = False
            else:
                self.handle_interact()

        if self.shop_open:
            self.update_shop()

        if pyxel.btnp(pyxel.KEY_M):
            pyxel.playm(0, loop=True)
        if pyxel.btnp(pyxel.KEY_N):
            pyxel.playm(1, loop=True)

        if self.message_timer > 0:
            self.message_timer -= dt
            if self.message_timer <= 0:
                self.message = ""

    # ---------------------------- TITLE ----------------------------

    def update_title(self, dt):
        if self.fade > 0:
            self.fade = max(0, self.fade - dt * 0.5)

        if pyxel.btnp(pyxel.KEY_UP):
            self.menu_index = max(0, self.menu_index - 1)
        if pyxel.btnp(pyxel.KEY_DOWN):
            self.menu_index = min(1, self.menu_index + 1)

        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            if self.menu_index == 0:
                self.state = "game"
            else:
                self.settings_index = 0
                self.state = "settings"

    # ---------------------------- TIMER MENU ----------------------------

    def update_timer_menu(self):
        options = [1, 5, 10, 15, 25, 45]
        if pyxel.btnp(pyxel.KEY_UP):
            self.timer_index = max(0, self.timer_index - 1)
        if pyxel.btnp(pyxel.KEY_DOWN):
            self.timer_index = min(len(options) - 1, self.timer_index + 1)
        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            self.start_timer(options[self.timer_index])

    def start_timer(self, minutes):
        reward = max(1, minutes // 3)
        self.reward_timer.start(minutes, reward)
        self.message = f"Timer: {minutes} min"
        self.message_timer = 3
        self.timer_menu_open = False

    # ---------------------------- INTERACTION ----------------------------

    def handle_interact(self):
        px, py = self.player.x, self.player.y

        # bookshelf
        if 64 < px < 96 and 32 < py < 80:
            self.message = random.choice([
                "Distant galaxies…",
                "A cat in Paris…",
                "Art of doing nothing…",
            ])
            self.message_timer = 3
            if random.random() < 0.6:
                self.coins += 1
            return

        # desk → shop
        if 32 < px < 64 and 90 < py < 130:
            self.shop_open = True
            self.message = "Shop opened."
            self.message_timer = 2
            return

        # bed → dream
        if 0 < px < 32 and 32 < py < 100:
            self.enter_dream()
            return

        # cat
        if self.cat.is_near(px, py):
            self.message = "The cat purrs softly."
            self.message_timer = 2
            if random.random() < 0.5:
                self.cat.start_follow()
            return

    def enter_dream(self):
        self.dream = DreamLevel(self.current_dream_level, self)
        self.dream.reset()
        self.state = "dream"
        self.message = ""

    # ---------------------------- SHOP ----------------------------

    def update_shop(self):
        if pyxel.btnp(pyxel.KEY_1):
            self.try_buy("Rug", 5)
        if pyxel.btnp(pyxel.KEY_2):
            self.try_buy("Fairy Lights", 8)
        if pyxel.btnp(pyxel.KEY_3):
            self.try_buy("Poster", 6)

    def try_buy(self, item, price):
        if self.coins < price:
            self.message = f"Need {price} coins!"
            self.message_timer = 2
            return
        self.coins -= price
        self.message = f"Bought {item}!"
        self.message_timer = 2

    # ---------------------------- DRAW ----------------------------

    def draw(self):
        skip_lighting = self.shop_open or self.timer_menu_open

        if self.state == "title":
            self.draw_title()
            return

        if self.state == "settings":
            self.draw_settings()
            return

        if self.state == "dream":
            self.dream.draw()
            return

        pyxel.cls(0)
        pyxel.bltm(0, 0, 0, 0, 0, SCREEN_W, SCREEN_H)

        if not skip_lighting and self.setting_lighting:
            self.draw_lighting()

        self.cat.draw()
        self.player.draw()
        self.draw_hud()
        self.reward_timer.draw()

        if self.shop_open:
            self.draw_shop()

        if self.timer_menu_open:
            self.draw_timer_menu()

    # ---------------------------- SETTINGS ----------------------------

    def update_settings(self, dt):
        num_options = 4  # added theme option

        if pyxel.btnp(pyxel.KEY_UP):
            self.settings_index = max(0, self.settings_index - 1)
        if pyxel.btnp(pyxel.KEY_DOWN):
            self.settings_index = min(num_options - 1, self.settings_index + 1)

        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            if self.settings_index == 0:
                self.setting_music = not self.setting_music
                if self.setting_music:
                    pyxel.playm(0, loop=True)
                else:
                    pyxel.stop()

            elif self.settings_index == 1:
                self.setting_lighting = not self.setting_lighting

            elif self.settings_index == 2:
                self.setting_speed = (self.setting_speed + 1) % 3

            elif self.settings_index == 3:
                self.setting_theme = (self.setting_theme + 1) % 3

        if pyxel.btnp(pyxel.KEY_Q) or pyxel.btnp(pyxel.KEY_BACKSPACE):
            self.state = "title"

    def draw_settings(self):
        pyxel.cls(0)
        pyxel.rect(0, 0, SCREEN_W, SCREEN_H, 0)

        pyxel.text(SCREEN_W // 2 - 16, 20, "SETTINGS", 10)
        pyxel.line(10, 30, SCREEN_W - 10, 30, 5)

        speed_labels = ["SLOW", "NORMAL", "FAST"]
        theme_labels = ["Classic", "Sunset", "Ocean"]

        options = [
            ("Music",    "ON" if self.setting_music else "OFF"),
            ("Lighting", "ON" if self.setting_lighting else "OFF"),
            ("Speed",    speed_labels[self.setting_speed]),
            ("Theme",    theme_labels[self.setting_theme]),
        ]

        for i, (label, value) in enumerate(options):
            y = 50 + i * 18
            is_sel = i == self.settings_index
            row_col = 10 if is_sel else 7
            prefix = "> " if is_sel else "  "
            val_col = 11 if value in ("ON", "NORMAL", "Classic") else 8
            pyxel.text(16, y, prefix + label, row_col)
            pyxel.text(90, y, value, val_col)

        pyxel.line(10, SCREEN_H - 22, SCREEN_W - 10, SCREEN_H - 22, 5)
        pyxel.text(10, SCREEN_H - 14, "SPC=toggle  Q=back", 5)

    # ---------------------------- TITLE SCREEN ----------------------------

    def draw_title(self):
        pyxel.cls(0)
        pyxel.bltm(0, 0, 0, 0, 0, SCREEN_W, SCREEN_H)
        pyxel.rect(0, 0, SCREEN_W, SCREEN_H, 0)

        title = "MIDNIGHT CORNER"
        pyxel.text(SCREEN_W // 2 - len(title) * 2, 50, title, 7)

        options = ["START", "SETTINGS"]
        for i, text in enumerate(options):
            color = 10 if i == self.menu_index else 7
            prefix = "> " if i == self.menu_index else "  "
            pyxel.text(SCREEN_W // 2 - 20, 90 + i * 12, prefix + text, color)

        if self.fade > 0:
            pyxel.rect(0, 0, SCREEN_W, SCREEN_H, 0)

    # ---------------------------- LIGHTING ----------------------------

    def draw_lighting(self):
        pyxel.rect(0, 0, SCREEN_W, SCREEN_H, 1)
        cx, cy = 100, 50
        for r, col in [(35, 5), (24, 6), (14, 7)]:
            pyxel.circ(cx, cy, r, col)
        wx, wy = 30, 30
        for r, col in [(28, 2), (18, 3)]:
            pyxel.circ(wx, wy, r, col)
        pyxel.bltm(0, 0, 0, 0, 0, SCREEN_W, SCREEN_H)

    # ---------------------------- HUD ----------------------------

    def draw_hud(self):
        pyxel.blt(4, 4, 0, 0, 128, 8, 8, 2)
        pyxel.text(14, 5, str(self.coins), 10)
        pyxel.text(4, 14, "SPACE:act T:timer", 7)
        pyxel.text(4, 22, "Bed=Dream!", 6)

        if self.message:
            w = len(self.message) * 4 + 4
            x = SCREEN_W // 2 - w // 2
            pyxel.rect(x - 2, 2, w + 4, 12, 1)
            pyxel.text(x, 4, self.message, 7)

    # ---------------------------- SHOP UI ----------------------------

    def draw_shop(self):
        x, y, w, h = 14, 55, 100, 75
        pyxel.rect(x, y, w, h, 1)
        pyxel.rectb(x, y, w, h, 7)
        pyxel.text(x + 6, y + 4, "SHOP", 10)
        pyxel.text(x + 6, y + 18, "1) Rug        (5)", 7)
        pyxel.text(x + 6, y + 28, "2) Fairy Lights(8)", 7)
        pyxel.text(x + 6, y + 38, "3) Poster     (6)", 7)
        pyxel.text(x + 6, y + 54, "SPACE to close", 5)

    # ---------------------------- TIMER MENU UI ----------------------------

    def draw_timer_menu(self):
        options = ["1 min", "5 min", "10 min", "15 min", "25 min", "45 min"]
        x, y, w, h = 10, 30, 108, 100
        pyxel.rect(x, y, w, h, 1)
        pyxel.rectb(x, y, w, h, 7)
        pyxel.text(x + 8, y + 6, "STUDY TIMER", 10)

        for i, label in enumerate(options):
            color = 10 if i == self.timer_index else 7
            prefix = "> " if i == self.timer_index else "  "
            pyxel.text(x + 8, y + 20 + i * 11, prefix + label, color)

        pyxel.text(x + 8, y + 88, "UP/DN+SPC  T=close", 5)


# ============================================================
#  MAIN ENTRY POINT
# ============================================================

if __name__ == "__main__":
    App()