import gradio as gr
import spaces
import numpy as np
import noise
import random
import json
import uuid
import os
import tempfile

from fastapi import FastAPI
from fastapi.responses import JSONResponse
# ---------- Генератор мира ----------
def generate_world(seed: int, size: int):
    np.random.seed(seed)
    random.seed(seed)
    scale = 30.0
    octaves = 6
    persistence = 0.5
    lacunarity = 2.0
    temp_noise = np.zeros((size, size))
    moist_noise = np.zeros((size, size))
    for x in range(size):
        for z in range(size):
            temp_noise[x][z] = noise.pnoise2(x/scale/2+seed, z/scale/2+seed, octaves=4)
            moist_noise[x][z] = noise.pnoise2(x/scale/2+seed*2, z/scale/2+seed*2, octaves=4)
    temp_noise = (temp_noise - temp_noise.min()) / (temp_noise.max() - temp_noise.min() + 1e-8)
    moist_noise = (moist_noise - moist_noise.min()) / (moist_noise.max() - moist_noise.min() + 1e-8)
    world_blocks = []
    for x in range(size):
        for z in range(size):
            h = noise.pnoise2(x/scale+seed, z/scale+seed, octaves=octaves,
                              persistence=persistence, lacunarity=lacunarity)
            c = noise.pnoise2(x/scale*2+seed*3, z/scale*2+seed*3, octaves=4)
            h = h * 0.7 + c * 0.3
            height_norm = (h + 1) / 2.0
            t = temp_noise[x][z]
            m = moist_noise[x][z]
            water_level = 0.45
            mountain_threshold = 0.75
            y_max = int(10 + height_norm * 15)
            for y in range(0, y_max + 1):
                if y == 0:
                    block = "bedrock"
                elif y <= 3:
                    block = "stone"
                elif y <= y_max - 4:
                    block = "dirt"
                else:
                    if height_norm < water_level:
                        block = "water" if y_max <= 12 else "sand"
                    elif height_norm > mountain_threshold:
                        block = "stone" if y == y_max else "dirt"
                    else:
                        if t < 0.3 and m < 0.3:
                            block = "sand" if y == y_max else "sandstone"
                        elif t < 0.3 and m > 0.6:
                            block = "grass" if y == y_max else "dirt"
                        elif t > 0.7 and m > 0.5:
                            block = "grass" if y == y_max else "dirt"
                            if y == y_max and random.random() < 0.1:
                                _make_tree(world_blocks, x, y, z)
                        else:
                            block = "grass" if y == y_max else "dirt"
                if block != "bedrock":
                    world_blocks.append({"x": x, "y": y, "z": z, "block": block})
    world_blocks.sort(key=lambda b: (b["y"], b["x"], b["z"]))
    return world_blocks

def _make_tree(blocks_list, x, y, z, trunk_height=3):
    for dy in range(1, trunk_height+1):
        blocks_list.append({"x": x, "y": y+dy, "z": z, "block": "wood"})
    for dx in (-1, 0, 1):
        for dz in (-1, 0, 1):
            for dy in range(trunk_height, trunk_height+2):
                if abs(dx) == 1 and abs(dz) == 1:
                    continue
                blocks_list.append({"x": x+dx, "y": y+dy, "z": z+dz, "block": "leaves"})

def chunk_world(blocks):
    full_json = json.dumps(blocks, ensure_ascii=False)
    max_len = 20000
    if len(full_json) <= max_len:
        return [full_json]
    parts = []
    start = 0
    while start < len(blocks):
        chunk_blocks = []
        current_len = 1
        idx = start
        while idx < len(blocks):
            block_str = json.dumps(blocks[idx], ensure_ascii=False)
            extra = 0 if not chunk_blocks else 1
            if current_len + len(block_str) + extra + 1 > max_len:
                break
            chunk_blocks.append(blocks[idx])
            current_len += len(block_str) + (1 if chunk_blocks else 0)
            idx += 1
        if not chunk_blocks:
            raise RuntimeError("Block too large for chunk")
        parts.append(json.dumps(chunk_blocks, ensure_ascii=False))
        start = idx
    return parts

# ---------- GPU только для генерации ----------
@spaces.GPU
def generate(seed: int, size: int):
    world = generate_world(seed, size)
    chunks = chunk_world(world)
    session_id = str(uuid.uuid4())
    filepath = os.path.join(tempfile.gettempdir(), f"mc_session_{session_id}.json")
    with open(filepath, "w", encoding="utf-8") as f:
        json.dump({"chunks": chunks, "total_parts": len(chunks), "seed": seed, "size": size}, f)
    return {"session_id": session_id, "total_parts": len(chunks)}

# ---------- Функция получения части (CPU, без очереди) ----------
def get_part(session_id: str, part_num: int):

    part_num = int(part_num)

    filepath = os.path.join(
        tempfile.gettempdir(),
        f"mc_session_{session_id}.json"
    )

    if not os.path.exists(filepath):
        return {
            "error": "Session not found"
        }

    with open(filepath, "r", encoding="utf-8") as f:
        data = json.load(f)

    if part_num < 0 or part_num >= data["total_parts"]:
        return {
            "error": "Invalid part number",
            "max": data["total_parts"]
        }

    return json.loads(
        data["chunks"][part_num]
    )


# ---------- Gradio UI ----------

with gr.Blocks() as demo:

    gr.Markdown(
        "# 🟫 Minecraft World Generator (ZeroGPU)"
    )

    with gr.Tab("Generate"):

        seed_in = gr.Number(
            value=42,
            label="Seed",
            precision=0
        )

        size_in = gr.Slider(
            minimum=20,
            maximum=50,
            value=35,
            step=5,
            label="Size"
        )

        gen_btn = gr.Button(
            "Generate World"
        )

        gen_out = gr.JSON(
            label="Result"
        )


    with gr.Tab("Get Part"):

        sess_in = gr.Textbox(
            label="Session ID"
        )

        part_in = gr.Number(
            value=0,
            label="Part Number",
            precision=0
        )

        get_btn = gr.Button(
            "Get Part"
        )

        part_out = gr.JSON(
            label="Chunk"
        )


    # GPU генерация
    gen_btn.click(
        fn=generate,
        inputs=[
            seed_in,
            size_in
        ],
        outputs=gen_out
    )


    # CPU получение блока
    get_btn.click(
        fn=get_part,
        inputs=[
            sess_in,
            part_in
        ],
        outputs=part_out
    )


app = demo.app

@app.get("/api/get_part")
def api_get_part(session_id: str, part_num: int):
    return JSONResponse(
        content=get_part(session_id, part_num)
    )

demo.launch()