import gradio as gr
import numpy as np
import random
from typing import Dict, List
import json

# Mock InstantMesh (replace with real model later)
def generate_mesh(chunk_x: int, chunk_z: int) -> List[List[float]]:
    chunk_size = 5
    mesh = np.zeros((chunk_size, chunk_size))
    prompt = "rocky hill" if random.random() > 0.5 else "flat plains"
    if "hill" in prompt:
        mesh[2, 2] = random.randint(2, 5)
        for i in range(chunk_size):
            for j in range(chunk_size):
                dist = np.sqrt((i - 2)**2 + (j - 2)**2)
                mesh[i, j] = max(0, 5 - dist)
    else:
        mesh = np.ones((chunk_size, chunk_size)) * random.uniform(0, 1)
    return mesh.tolist()

# Game state
class Game:
    def __init__(self):
        self.player_pos = [0, 0]
        self.world: Dict[tuple, List[List[float]]] = {}
        self.chunk_size = 5

    def get_chunk_coords(self, x: int, z: int) -> tuple:
        return (x // self.chunk_size, z // self.chunk_size)

    def generate_chunk(self, chunk_x: int, chunk_z: int):
        key = (chunk_x, chunk_z)
        if key not in self.world:
            print(f"Generating chunk at {key}")
            self.world[key] = generate_mesh(chunk_x, chunk_z)
        return self.world[key]

    def move(self, dx: int, dz: int):
        self.player_pos[0] += dx
        self.player_pos[1] += dz
        chunk_x, chunk_z = self.get_chunk_coords(self.player_pos[0], self.player_pos[1])
        chunk = self.generate_chunk(chunk_x, chunk_z)
        return {
            "player_pos": self.player_pos,
            "chunk": chunk,
            "chunk_coords": [chunk_x, chunk_z]
        }

game = Game()

# Gradio interface with custom endpoint simulation
with gr.Blocks() as app:
    state = gr.State(value=game)

    def move_player(dx, dz, game_state):
        result = game_state.move(dx, dz)
        return game_state, json.dumps(result)

    # Hidden button to simulate API call
    move_btn = gr.Button("Move", visible=False)
    output = gr.JSON(label="Game State")

    # Trigger move on input
    dx_input = gr.Number(label="DX", value=0, visible=False)
    dz_input = gr.Number(label="DZ", value=0, visible=False)
    move_btn.click(
        fn=move_player,
        inputs=[dx_input, dz_input, state],
        outputs=[state, output]
    )

    # Custom route for Three.js to fetch from
    def get_game_state(dx: int = 0, dz: int = 0):
        result = game.move(dx, dz)
        return result

    # Expose a pseudo-API endpoint (not standard Gradio, but works with Spaces)
    app.queue()
    app.launch(share=True)  # Remove share=True when deploying to Spaces

# For Hugging Face Spaces, ensure this is the entry point
if __name__ == "__main__":
    app.launch(server_port=7860)