from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from llama_cpp import Llama
import uvicorn
import re

app = FastAPI()

# Standard CORS setup for local web tools
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)

# Load the fine-tuned math model
# Utilizing the full 16GB VRAM of the RTX 4060 Ti
llm = Llama(
    model_path="/home/khurram/ai_models/math_dataset/math_viz_Q4_K_M.gguf",
    n_gpu_layers=-1, 
    n_ctx=2048,
    n_batch=512,
    temperature=0.1
)

def clean_and_format_ggb(raw_text):
    """
    Standardizes model coordinates and brackets for GeoGebra.
    Converts [Cylinder[<0,0,0>, <3,0,0>, <0,10,0>]] 
    to Cylinder((0,0,0), (0,10,0), 3.0)
    """
    # 1. Clean up bracket variations
    text = raw_text.replace("<", "(").replace(">", ")")
    
    # 2. Extract all coordinate sets (x,y,z)
    coords = re.findall(r"\((-?\d+\.?\d*,\s*-?\d+\.?\d*,\s*-?\d+\.?\d*)\)", text)
    
    if len(coords) >= 3:
        bottom_pt = f"({coords[0]})"
        top_pt = f"({coords[2]})"
        
        # Extract scalar radius from the middle point
        radius_match = re.findall(r"[-+]?\d*\.\d+|\d+", coords[1])
        radius = next((abs(float(n)) for n in radius_match if float(n) != 0), 3.0)
        
        return f"Cylinder({bottom_pt}, {top_pt}, {radius})"
    
    # Fallback for simple Sphere or direct commands
    return text.replace("[", "").replace("]", "").replace("<", "(").replace(">", ")").strip()

@app.post("/ask")
async def ask_geo(data: dict):
    user_prompt = data.get("prompt", "")
    
    # Step 1: Request the "Thought" (Mathematical Reasoning)
    prompt = f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>thought\n"
    thought_output = llm(prompt, max_tokens=150, stop=["<|im_end|>"])
    thought_text = thought_output["choices"][0]["text"].strip()

    # Step 2: Request the "Assistant" (GeoGebra Code)
    command_prompt = f"{prompt}{thought_text}<|im_end|>\n<|im_start|>assistant\n"
    command_output = llm(command_prompt, max_tokens=150, stop=["<|im_end|>"])
    assistant_raw = command_output["choices"][0]["text"].strip()

    # Final string formatting for the GeoGebra Applet
    final_cmds = clean_and_format_ggb(assistant_raw)
    
    return {
        "commands": final_cmds, 
        "thought": thought_text
    }

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)