import pygame
import os
import sys
import contextlib

class Window:
    # friendly string -> pygame constant mapping
    KEY_MAP = {
        "left": pygame.K_LEFT,
        "right": pygame.K_RIGHT,
        "up": pygame.K_UP,
        "down": pygame.K_DOWN,
        "space": pygame.K_SPACE,
        "enter": pygame.K_RETURN,
        "backspace": pygame.K_BACKSPACE,
        "tab": pygame.K_TAB,
        "escape": pygame.K_ESCAPE,
        "Rshift": pygame.K_LSHIFT,
        "Lctrl": pygame.K_LCTRL,
        "Rshift": pygame.K_RSHIFT,
        "Rctrl": pygame.K_RCTRL,
        "alt": pygame.K_LALT,
        "keya": pygame.K_a,
        "keyb": pygame.K_b,
        "keyc": pygame.K_c,
        "keyd": pygame.K_d,
        "keye": pygame.K_e,
        "keyf": pygame.K_f,
        "keyg": pygame.K_g,
        "keyh": pygame.K_h,
        "keyi": pygame.K_i,
        "keyj": pygame.K_j,
        "keyk": pygame.K_k,
        "keyl": pygame.K_l,
        "keym": pygame.K_m,
        "keyn": pygame.K_n,
        "keyo": pygame.K_o,
        "keyp": pygame.K_p,
        "keyq": pygame.K_q,
        "keyr": pygame.K_r,
        "keys": pygame.K_s,
        "keyt": pygame.K_t,
        "keyu": pygame.K_u,
        "keyv": pygame.K_v,
        "keyw": pygame.K_w,
        "keyx": pygame.K_x,
        "keyy": pygame.K_y,
        "keyz": pygame.K_z,
        "0": pygame.K_0,
        "1": pygame.K_1,
        "2": pygame.K_2,
        "3": pygame.K_3,
        "4": pygame.K_4,
        "5": pygame.K_5,
        "6": pygame.K_6,
        "7": pygame.K_7,
        "8": pygame.K_8,
        "9": pygame.K_9,
        "f1": pygame.K_F1,
        "f2": pygame.K_F2,
        "f3": pygame.K_F3,
        "f4": pygame.K_F4,
        "f5": pygame.K_F5,
        "f6": pygame.K_F6,
        "f7": pygame.K_F7,
        "f8": pygame.K_F8,
        "f9": pygame.K_F9,
        "f10": pygame.K_F10,
        "f11": pygame.K_F11,
        "f12": pygame.K_F12,
        "insert": pygame.K_INSERT,
        "delete": pygame.K_DELETE,
        "home": pygame.K_HOME,
        "end": pygame.K_END,
        "pageup": pygame.K_PAGEUP,
        "pagedown": pygame.K_PAGEDOWN,
        "capslock": pygame.K_CAPSLOCK,
        "numlock": pygame.K_NUMLOCK,
        "scrolllock": pygame.K_SCROLLOCK,
    }


    def __init__(self, width=640, height=480, title="Game Window", icon=None):
        # temporarily silence pygame print
        with open(os.devnull, "w") as f, contextlib.redirect_stdout(f):
            pygame.init()  # <- no more hello message

        self.width = width
        self.height = height
        self.title = title
        self.screen = pygame.display.set_mode((width, height))
        pygame.display.set_caption(title)

        # --- ICON HANDLING ---
        if icon:
            # user-provided icon logic
            if isinstance(icon, str):
                icon_path = icon
                if not os.path.isabs(icon_path):
                    icon_path = os.path.join(os.getcwd(), icon)
                if os.path.exists(icon_path):
                    img = pygame.image.load(icon_path).convert_alpha()
                    pygame.display.set_icon(img)
                else:
                    print(f"⚠️ Icon file not found: {icon_path}")
            else:
                pygame.display.set_icon(icon)
        else:
            default_icon_path = os.path.join(os.path.dirname(__file__), "icons", "default_icon.png")
            if os.path.exists(default_icon_path):
                img = pygame.image.load(default_icon_path).convert_alpha()
                pygame.display.set_icon(img)
            else:
                default_icon = pygame.Surface((32,32))
                default_icon.fill((0,255,0))
                pygame.draw.rect(default_icon, (0,200,0), (8,8,16,16))
                pygame.display.set_icon(default_icon)

        self.clock = pygame.time.Clock()
        self.running = True
        self.keys = {}
        self.mouse_pos = (0,0)
        self.mouse_pressed = (False, False, False)
        self.draw_calls = []

    # === input helpers ===
    def is_key_pressed(self, key):
        # key can be pygame constant OR string from KEY_MAP
        if isinstance(key, str):
            key = self.KEY_MAP.get(key, None)
            if key is None:
                return False
        return self.keys.get(key, False)

    def is_mouse_pressed(self, button=1):
        return self.mouse_pressed[button-1]

    # === font system ===
    def get_font(self, name="Arial", size=16, bold=False, italic=False):
        if not hasattr(self, "font_cache"):
            self.font_cache = {}
        key = (name, size, bold, italic)
        if key not in self.font_cache:
            try:
                font = pygame.font.SysFont(name, size, bold, italic)
            except:
                font = pygame.font.Font(None, size)
            self.font_cache[key] = font
        return self.font_cache[key]

    def draw_text(self, text, font="Arial", size=16, color=(255,255,255), pos=(0,0), bold=False, italic=False):
        f = self.get_font(font, size, bold, italic)
        surf = f.render(text, True, color)
        self.screen.blit(surf, pos)

    # === drawing helpers ===a
    def draw_rect(self, color, rect):
        self.draw_calls.append(("rect", color, rect))

    def draw_circle(self, color, pos, radius):
        self.draw_calls.append(("circle", color, pos, radius))

    def draw_sprite(self, sprite):
        self.draw_calls.append(("sprite", sprite))

    # === main loop ===
    def run(self, update_func=None):
        while self.running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                elif event.type == pygame.KEYDOWN:
                    self.keys[event.key] = True
                elif event.type == pygame.KEYUP:
                    self.keys[event.key] = False
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    self.mouse_pressed = pygame.mouse.get_pressed()
                elif event.type == pygame.MOUSEBUTTONUP:
                    self.mouse_pressed = pygame.mouse.get_pressed()
                elif event.type == pygame.MOUSEMOTION:
                    self.mouse_pos = event.pos

            if update_func:
                update_func(self)

            self.screen.fill((20,20,50))
            for call in self.draw_calls:
                if call[0] == "rect":
                    pygame.draw.rect(self.screen, call[1], call[2])
                elif call[0] == "circle":
                    pygame.draw.circle(self.screen, call[1], call[2], call[3])
                elif call[0] == "sprite":
                    call[1].draw(self.screen)
            self.draw_calls.clear()

            pygame.display.flip()
            self.clock.tick(60)

        pygame.quit()
