# Neon Snake — Gradio Edition
# Works on Hugging Face Spaces (Gradio Space)
# Controls: Arrow buttons or WASD in the small textbox (press Enter)

import random
from typing import Dict, List, Tuple
from PIL import Image, ImageDraw, ImageFont
import gradio as gr

# ---------- Visual + Game Constants ----------
GRID = 24            # number of cells in each dimension
CELL = 18            # pixel size of a cell
BORDER = 12          # outer border size
W = H = GRID * CELL + BORDER * 2

NEON_BG  = (11, 15, 26)
GRID_COL = (35, 60, 139)
SNAKE    = (95, 168, 255)
HEAD     = (255, 255, 255)
FOOD     = (139, 233, 253)
TXT      = (167, 195, 255)

START_SPEED = 0.14   # seconds per tick (lower = faster)
SPEED_STEP  = 0.01   # how much to speed up each fruit
MIN_SPEED   = 0.06   # cap so it doesn't get too fast

DIRS = {
    "UP":    (0, -1),
    "DOWN":  (0,  1),
    "LEFT":  (-1, 0),
    "RIGHT": (1,  0),
}

OPPOSITE = {
    "UP": "DOWN", "DOWN": "UP", "LEFT": "RIGHT", "RIGHT": "LEFT"
}

# ---------- Game Logic ----------
def new_game() -> Dict:
    snake = [(GRID // 2, GRID // 2)]
    direction = "RIGHT"
    food = place_food(snake)
    return {
        "snake": snake,
        "direction": direction,
        "pending_dir": direction,  # queued turn to apply next tick
        "food": food,
        "alive": True,
        "score": 0,
        "high": 0,
        "paused": False,
        "speed": START_SPEED,
    }

def place_food(snake: List[Tuple[int, int]]) -> Tuple[int, int]:
    free = [(x, y) for x in range(GRID) for y in range(GRID) if (x, y) not in snake]
    return random.choice(free) if free else (0, 0)

def step(state: Dict) -> Dict:
    if not state["alive"] or state["paused"]:
        return state

    # apply any pending turn (disallow reversal into self)
    if state["pending_dir"] != OPPOSITE[state["direction"]]:
        state["direction"] = state["pending_dir"]

    dx, dy = DIRS[state["direction"]]
    head_x, head_y = state["snake"][0]
    new_head = (head_x + dx, head_y + dy)

    # collide with walls
    if not (0 <= new_head[0] < GRID and 0 <= new_head[1] < GRID):
        state["alive"] = False
        state["high"] = max(state["high"], state["score"])
        return state

    # collide with self
    if new_head in state["snake"]:
        state["alive"] = False
        state["high"] = max(state["high"], state["score"])
        return state

    # move
    state["snake"].insert(0, new_head)

    # eat?
    if new_head == state["food"]:
        state["score"] += 1
        state["food"] = place_food(state["snake"])
        state["speed"] = max(MIN_SPEED, state["speed"] - SPEED_STEP)
    else:
        state["snake"].pop()

    return state

# ---------- Rendering ----------
def render(state: Dict) -> Image.Image:
    img = Image.new("RGB", (W, H), NEON_BG)
    draw = ImageDraw.Draw(img)

    # grid
    for i in range(GRID + 1):
        x = BORDER + i * CELL
        y = BORDER + i * CELL
        draw.line([(BORDER, y), (W - BORDER, y)], fill=GRID_COL, width=1)
        draw.line([(x, BORDER), (x, H - BORDER)], fill=GRID_COL, width=1)

    # food
    fx, fy = state["food"]
    draw_cell(draw, fx, fy, FOOD, radius=6)

    # snake
    for i, (sx, sy) in enumerate(state["snake"]):
        color = HEAD if i == 0 else SNAKE
        draw_cell(draw, sx, sy, color, radius=6)

    # HUD
    status = "PAUSED" if state["paused"] else ("GAME OVER" if not state["alive"] else "NEON SNAKE")
    hud = f"{status}  |  Score: {state['score']}  High: {state['high']}"
    draw.text((BORDER, 4), hud, fill=TXT)

    return img

def draw_cell(draw: ImageDraw.ImageDraw, cx: int, cy: int, col, radius=4):
    x0 = BORDER + cx * CELL
    y0 = BORDER + cy * CELL
    x1 = x0 + CELL
    y1 = y0 + CELL
    # rounded rectangle
    draw.rounded_rectangle([x0 + 2, y0 + 2, x1 - 2, y1 - 2], radius=radius, fill=col)

# ---------- UI Callbacks ----------
def on_key(state: Dict, key: str):
    if not state["alive"]:
        return state
    k = (key or "").strip().upper()
    # WASD mapping
    if k in ["W", "UP"]:
        nd = "UP"
    elif k in ["S", "DOWN"]:
        nd = "DOWN"
    elif k in ["A", "LEFT"]:
        nd = "LEFT"
    elif k in ["D", "RIGHT"]:
        nd = "RIGHT"
    else:
        return state
    # don't allow immediate reversal into self
    if nd != OPPOSITE[state["direction"]]:
        state["pending_dir"] = nd
    return state

def click_dir(state: Dict, nd: str):
    # same behavior as keyboard
    if not state["alive"]:
        return state
    if nd != OPPOSITE[state["direction"]]:
        state["pending_dir"] = nd
    return state

def toggle_pause(state: Dict):
    if state["alive"]:
        state["paused"] = not state["paused"]
    return state

def reset_game(state: Dict):
    high = state.get("high", 0)
    ng = new_game()
    ng["high"] = max(high, ng["high"])
    return ng

def tick(state: Dict):
    # advance game and return new frame + (possibly) updated timer interval
    before = state["speed"]
    state = step(state)
    frame = render(state)
    # if speed changed after eating, we can notify the frontend to adjust timer
    return frame, state["speed"], state

# ---------- Gradio App ----------
with gr.Blocks(css="""
#wrap {max-width: 760px; margin: 0 auto;}
.button-row button {min-width: 84px;}
""") as demo:
    gr.Markdown("## 🐍 Neon Snake — Gradio\nUse arrow buttons or WASD (enter) • Eat food • Don’t hit walls or yourself")

    gstate = gr.State(new_game())

    with gr.Row(elem_id="wrap"):
        canvas = gr.Image(value=render(gstate.value), label="Game", width=W, height=H, show_download_button=False)

    with gr.Row(equal_height=True, elem_classes="button-row"):
        up = gr.Button("▲ Up")
        left = gr.Button("◀ Left")
        pause = gr.Button("⏯ Pause/Resume")
        right = gr.Button("Right ▶")
        down = gr.Button("Down ▼")
        reset = gr.Button("🔄 Reset")

    kb = gr.Textbox(label="WASD / Arrow words (e.g., 'up') then Enter", value="", placeholder="w / a / s / d or up / down / left / right")
    kb_cb = gr.Button("Apply Key")

    # Timer drives the game loop; interval updates dynamically as speed changes
    timer = gr.Timer(interval=START_SPEED, active=True)

    # --- Wire up controls ---
    up.click(click_dir, [gstate], [gstate], js=None, preprocess=True, postprocess=False, kwargs={"nd": "UP"})
    down.click(click_dir, [gstate], [gstate], kwargs={"nd": "DOWN"})
    left.click(click_dir, [gstate], [gstate], kwargs={"nd": "LEFT"})
    right.click(click_dir, [gstate], [gstate], kwargs={"nd": "RIGHT"})

    pause.click(toggle_pause, [gstate], [gstate])
    reset.click(reset_game, [gstate], [gstate]).then(lambda s: render(s), [gstate], [canvas])

    # Keyboard textbox
    def kb_apply(s: Dict, k: str):
        s = on_key(s, k)
        return "", s  # clear the textbox for the next key
    kb_cb.click(kb_apply, [gstate, kb], [kb, gstate])

    # Game loop
    def game_tick(state: Dict):
        frame, new_interval, new_state = tick(state)
        # return: image, timer interval, updated state
        return frame, new_interval, new_state

    timer.tick(game_tick, [gstate], [canvas, timer, gstate])

demo.queue().launch()