import json

import requests
from fastapi import FastAPI
from pydantic import BaseModel

from config import DEVICE, CHECKPOINT_DIR
from data.tokenizer import Tokenizer
from makescript import json_string_to_py
from model.builder_model import BuilderModel
from utils.checkpoint import latest_checkpoint, load_checkpoint

import subprocess
import sys
import os
import base64

app = FastAPI(title="Minecraft Builder AI")

if DEVICE == "cuda":
    print("Cuda Mode: true")
else:
    print("Cuda Mode: false")

print("Loading tokenizer...")
tokenizer = Tokenizer()
tokenizer.load("data")

print("Loading model...")
model = BuilderModel(tokenizer).to(DEVICE)

checkpoint = latest_checkpoint(CHECKPOINT_DIR)
load_checkpoint(
    checkpoint,
    model,
    optimizer=None,
    map_location=DEVICE,
)

model.eval()


class BuildRequest(BaseModel):
    prompt: str


@app.post("/generate")
def generate(req: BuildRequest):

    blocks, tokens = model.generate(req.prompt)

    return {
        "prompt": req.prompt,
        "tokens": tokens,
        "blocks": blocks,
    }

class RunRequest(BaseModel):
    script: str

@app.post("/run_script")
def runscript(req: RunRequest):
    result = subprocess.run([
        sys.executable,
        "builders/" + req.script + ".py"
    ])

    if result.returncode != 0:
        return {
            "success": False,
            "schem": None
        }

    schem_path = "schematics/" + req.script + ".schem"

    if not os.path.exists(schem_path):
        return {
            "success": False,
            "schem": None
        }

    with open(schem_path, "rb") as f:
        schem = base64.b64encode(f.read()).decode("ascii")

    return {
        "success": True,
        "schem": schem
    }

class AIScriptRequest(BaseModel):
    url: str
    model: str
    api_key: str
    prompt: str
    context_length: int
    temperature: float
    reasoning: str

@app.post("/aiscript")
def aiscript(req: AIScriptRequest):
    systemprompt = r"""
    You are a Minecraft Building AI.
    You build with the mcschematic python library.
    Your task is to return a valid script in the following format:
    {
        "file_name": FILE_NAME,
        "script: SCRIPT,
        "response": RESPONSE
    }
    
    Rules:
    - Only use valid Minecraft blocks
    - The output folder is called "schematics"
    - The schematic name must match with the FILE_NAME
    - In settings there always needs to be a variable called "SEED"
    
    Documentation:
    mcschematic is a Python library for programmatically creating Minecraft Java Edition schematic files. The basic workflow is:
    Create an MCSchematic object.
    Place Minecraft blocks at (x, y, z) coordinates with setBlock().
    Optionally manipulate blocks/entities/data.
    Save the resulting structure as a .schem file with save().
    The core conceptual model is:
    import mcschematic
    schematic = mcschematic.MCSchematic()
    schematic.setBlock((x, y, z), "minecraft:oak_log")
    schematic.setBlock((x, y, z), "minecraft:oak_leaves")
    schematic.save(
        "output_directory",
        "my_structure",
        mcschematic.Version.JE_1_20_1
    )
    Important concepts
    MCSchematic
    This is the main container representing the Minecraft structure.
    schem = mcschematic.MCSchematic()
    Think of it as an in-memory 3D voxel structure.
    Coordinates
    Blocks are addressed using:
    (x, y, z)
    For procedural generation, treat these as a normal 3D Cartesian coordinate system. Pick a clear origin convention—for example, (0, 0, 0) as the bottom-center of a tree—and build everything relative to that.
    Placing blocks
    The primary operation is:
    schem.setBlock((x, y, z), "minecraft:oak_log")
    Block IDs should use Minecraft's namespaced IDs:
    minecraft:stone
    minecraft:dirt
    minecraft:oak_log
    minecraft:oak_leaves
    minecraft:glass
    Block states can be specified as part of the block string when supported, e.g.:
    "minecraft:oak_leaves[persistent=true]"
    For a procedural generator, create helper functions around setBlock() rather than scattering raw calls throughout the code.
    Example:
    def block(schem, x, y, z, block_id):
    schem.setBlock((x, y, z), block_id)
    Then your generation code can operate at a higher level:
    block(schem, x, y, z, "minecraft:oak_log")
    Procedural generation
    mcschematic isn't a tree/terrain generator itself. It provides the Minecraft schematic representation and file output.
    The procedural geometry should be implemented by the agent.
    For example:
    Tree generator
    │
    ├── trunk()
    │   └── places log blocks
    │
    ├── branch()
    │   └── calculates a 3D line and places logs
    │
    ├── roots()
    │   └── creates radial structures
    │
    └── canopy()
    └── creates randomized leaf volumes
    │
    └── setBlock()
    Useful geometric primitives to implement yourself include:
    line between two 3D points
    cylinder
    sphere / ellipsoid
    cone
    tapered branch
    radial distribution
    noise/randomized volume
    distance-based falloff
    For example, a spherical leaf cluster can be generated by checking:
    distance = (
        (x - cx) ** 2 +
        (y - cy) ** 2 +
        (z - cz) ** 2
    )
    if distance <= radius ** 2:
        schem.setBlock(
            (x, y, z),
            "minecraft:oak_leaves[persistent=true]"
    )
    Minecraft version
    The library has Minecraft-version constants under mcschematic.Version. The agent should inspect the installed library rather than hard-code a version based on memory.
    For example:
    print(dir(mcschematic.Version))
    and choose the appropriate Java Edition version available in the installed package.
    Likewise, the agent should inspect the actual installed API:
    import mcschematic
    print(mcschematic.version)
    print(dir(mcschematic))
    If there is uncertainty about an API call, inspect the package source/signatures instead of inventing an API.
    Saving
    The schematic is eventually serialized with:
    schem.save(
        output_folder,
        schematic_name,
        version
    )
    The resulting file is intended to be usable by Minecraft schematic-compatible tools.
    Recommended architecture for a tree generator
    A good design would be:
    tree = TreeGenerator(
        height=20,
        trunk_radius=2,
        seed=1234,
    )
    tree.generate()
    tree.export(
        "output/tree.schem"
    )
    Internally:
    TreeGenerator
    ↓
    generate_trunk()
    ↓
    generate_branches()
    ↓
    generate_roots()
    ↓
    generate_canopy()
    ↓
    MCSchematic.setBlock()
    ↓
    MCSchematic.save()
    
    Helper Methods:
    def set_block(x, y, z, block):
        schem.setBlock((x, y, z), block)
        
    def fill_sphere(cx, cy, cz, radius, block, irregularity=0.0):
        # Create a branch between two points.
        for x in range(cx - radius, cx + radius + 1):
            for y in range(cy - radius, cy + radius + 1):
                for z in range(cz - radius, cz + radius + 1):
    
                    distance = math.sqrt(
                        (x - cx) ** 2 +
                        (y - cy) ** 2 +
                        (z - cz) ** 2
                    )

                    # Make the sphere slightly irregular.
                    wobble = random.uniform(-irregularity, irregularity)

                    if distance <= radius + wobble:
                        set_block(x, y, z, block)


    def cylinder(cx, cz, y_start, y_end, radius, block):
        # Create a vertical cylinder.
            for y in range(y_start, y_end + 1):
                for x in range(cx - radius, cx + radius + 1):
                    for z in range(cz - radius, cz + radius + 1):

                        if (x - cx) ** 2 + (z - cz) ** 2 <= radius ** 2:
                            set_block(x, y, z, block)


    def branch(x1, y1, z1, x2, y2, z2, radius=1):
        # Create a branch between two points.
            steps = max(
                abs(x2 - x1),
                abs(y2 - y1),
                abs(z2 - z1)
            )

            for i in range(steps + 1):
                t = i / max(steps, 1)

                x = round(x1 + (x2 - x1) * t)
                y = round(y1 + (y2 - y1) * t)
                z = round(z1 + (z2 - z1) * t)

                cylinder(x, z, y, y, radius, "minecraft:oak_log")
                
    These are example helper methods. Also try to make your own when needed.

    This is how the schematic will be saved:
        schem.save(
        OUTPUT_FOLDER,
        OUTPUT_NAME,
        MC_VERSION
    )

    print(f"Saved {OUTPUT_FOLDER}/{OUTPUT_NAME}.schem")
    
    Coding Rules ONLY APPLY THESE IN THE SCRIPT NOT THE ENTIRE JSON:
    file_name:
    Use only letters, numbers, underscores, and hyphens.
    Do not include .py.
    Do not include directories or paths.
    script:
    Must contain the complete, executable Python source code.
    The script must use mcschematic.
    The script must save the generated schematic using schem.save(...).
    Use Minecraft Java Edition 1.20.1 unless the user specifies another version.
    The script must be self-contained and must not require the user to manually modify it.
    CRITICAL JSON REQUIREMENT:
    The value of script is a JSON string.
    You MUST properly JSON-escape the Python source code.
    Every double quote (") INSIDE the Python script must be escaped as \".
    Every newline inside the Python script must be represented as \n.
    Never output literal unescaped newlines inside the JSON script string.
    Never produce invalid JSON.
    Do NOT wrap the response in Markdown code fences.
    Do NOT output json. Do NOT output python.
    Output ONLY the JSON object.
    Do not add explanations, comments outside the Python script, or any text before or after the JSON object.
    Before responding, internally verify that your entire response is valid JSON and could be successfully parsed with Python's json.loads().
    The resulting value of:
    json.loads(response)["script"]
    must be directly writable to a .py file and must produce valid Python source code.
    In "response" give a short description what you build.
    Example of the REQUIRED output format:
    {
        "file_name": "small_oak_tree",
        "script": "import mcschematic\nimport math\n\nMC_VERSION = mcschematic.Version.JE_1_20_1\nschem = mcschematic.MCSchematic()\n\nschem.setBlock((0, 0, 0), \"minecraft\")\n\nschem.save(\"schematics\", \"small_oak_tree\", MC_VERSION)\n",
        "response": "I build your small oak tree."
    }
    Remember: the example above demonstrates the JSON encoding requirements. Your actual response must contain ONLY one valid JSON object.
    """

    payload = {
        "model": req.model,
        "messages": [
            {
                "role": "system",
                "content": systemprompt
            },
            {
                "role": "user",
                "content": req.prompt
            }
        ],
        "max_completion_tokens": req.context_length,
        "reasoning_effort": req.reasoning,
        "temperature": req.temperature,
        "stream": True,
    }

    headers = {
        "Authorization": f"Bearer {req.api_key}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(
            req.url + "/chat/completions",
            json=payload,
            headers=headers,
            timeout=600,
            stream=True
        )
        response.raise_for_status()

        reasoning = ""
        answer = ""

        for line in response.iter_lines(decode_unicode=True):
            if not line or not line.startswith("data: "):
                continue

            data = line[6:]

            if data == "[DONE]":
                break

            try:
                chunk = json.loads(data)
            except json.JSONDecodeError:
                print("Could not decode SSE chunk:", data)
                continue

            # Some chunks contain no choices.
            choices = chunk.get("choices", [])
            if not choices:
                continue

            choice = choices[0]
            delta = choice.get("delta", {})

            # Reasoning tokens
            reasoning_token = delta.get("reasoning_content")
            if reasoning_token:
                reasoning += reasoning_token
                print(reasoning_token, end="", flush=True)

            # Normal answer tokens
            content_token = delta.get("content")
            if content_token:
                answer += content_token
                print(content_token, end="", flush=True)

        print("\n\nFINAL ANSWER:")
        print(answer)

        # Parse the AI's JSON
        try:
            result = json.loads(answer)
        except json.JSONDecodeError as e:
            print("AI did not return valid JSON!")
            print("JSON error:", e)
            print("Raw answer:")
            print(repr(answer))

            return {
                "success": False,
                "file_name": "failed",
                "response": "AI returned invalid JSON"
            }

        # Validate expected fields
        if "file_name" not in result:
            raise ValueError("AI JSON is missing 'file_name'")

        if "script" not in result:
            raise ValueError("AI JSON is missing 'script'")

        if "response" not in result:
            raise ValueError("AI JSON is missing 'response'")

        file_name = result["file_name"]
        script = result["script"]
        remessage = result["response"]

        # Write the actual Python script
        with open(file_name + ".py", "w", encoding="utf-8") as f:
            f.write(script)

        return {
            "success": True,
            "file_name": file_name,
            "response": remessage
        }

    except requests.exceptions.RequestException as e:
        print("API request failed!")
        print("Error:", e)

        if e.response is not None:
            print("Status code:", e.response.status_code)
            print("Response:", e.response.text)

        return {
            "success": False,
            "file_name": "failed",
            "response": "failed"
        }

    except Exception as e:
        print("Unexpected error:", e)

        return {
            "success": False,
            "file_name": "failed",
            "response": str(e)
        }