import os
import json
import shutil
import time
import threading

from gradio import Server
from gradio_client import Client
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles

# ---------------------------------------------------------------------------
# Persistent storage
# ---------------------------------------------------------------------------
DATA_DIR = os.environ.get("DATA_DIR", "/data")
GRID_FILE = os.path.join(DATA_DIR, "grid.json")
IMAGES_DIR = os.path.join(DATA_DIR, "images")

os.makedirs(IMAGES_DIR, exist_ok=True)

_grid_lock = threading.Lock()


def load_grid():
    with _grid_lock:
        if os.path.exists(GRID_FILE):
            with open(GRID_FILE) as f:
                return json.load(f)
    return {}


def save_grid(grid):
    with _grid_lock:
        with open(GRID_FILE, "w") as f:
            json.dump(grid, f)


# ---------------------------------------------------------------------------
# FLUX Klein client (lazy init — Space may be sleeping)
# ---------------------------------------------------------------------------
_flux = None


def get_flux():
    global _flux
    if _flux is None:
        _flux = Client("black-forest-labs/flux-klein-9b-kv")
    return _flux


# ---------------------------------------------------------------------------
# Server
# ---------------------------------------------------------------------------
app = Server()

# Serve generated images as static files
app.mount("/images", StaticFiles(directory=IMAGES_DIR), name="images")


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


@app.api(name="get_grid")
def get_grid() -> str:
    """Return all placed buildings as JSON."""
    return json.dumps(load_grid())


@app.api(name="place_building")
def place_building(x: int, y: int, prompt: str) -> str:
    """Generate an isometric building and place it on the grid."""
    key = f"{x},{y}"
    grid = load_grid()

    if key in grid:
        return json.dumps({"error": "Tile already occupied"})

    full_prompt = (
        f"isometric building, white background, game asset, "
        f"single building centered, no ground, no shadow, clean edges, "
        f"{prompt}"
    )

    try:
        result, seed = get_flux().predict(
            prompt=full_prompt,
            input_images=[],
            seed=0,
            randomize_seed=True,
            width=512,
            height=512,
            num_inference_steps=4,
            prompt_upsampling=False,
            api_name="/generate",
        )
    except Exception as e:
        return json.dumps({"error": str(e)})

    # Copy generated image to persistent storage
    filename = f"{x}_{y}.webp"
    dest = os.path.join(IMAGES_DIR, filename)
    shutil.copy2(result, dest)

    grid[key] = {
        "prompt": prompt,
        "seed": seed,
        "ts": int(time.time()),
    }
    save_grid(grid)

    return json.dumps({"ok": True, "image": f"/images/{filename}", "seed": seed})


@app.api(name="remove_building")
def remove_building(x: int, y: int) -> str:
    """Remove a building from the grid."""
    key = f"{x},{y}"
    grid = load_grid()

    if key not in grid:
        return json.dumps({"error": "No building at this tile"})

    del grid[key]
    save_grid(grid)

    img_path = os.path.join(IMAGES_DIR, f"{x}_{y}.webp")
    if os.path.exists(img_path):
        os.remove(img_path)

    return json.dumps({"ok": True})


demo = app

if __name__ == "__main__":
    demo.launch(show_error=True, ssr_mode=False)