import os
import json
from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

app = FastAPI()

MODEL_ID = "Qwen/Qwen2.5-Math-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# --- SMART GPU/CPU FALLBACK ---
has_gpu = torch.cuda.is_available()
device_target = "cuda" if has_gpu else "cpu"
dtype_target = torch.float16 if has_gpu else torch.bfloat16

print(f"🚀 Booting Local AI on: {device_target.upper()} with {dtype_target}")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, 
    torch_dtype=dtype_target,
    low_cpu_mem_usage=True,
    device_map="auto" if has_gpu else None
)

if not has_gpu:
    model = model.to("cpu")

with open("GUIDE-STRUCTURE.md", "r") as f:
    guide_prompt = f.read()

with open("EXPLAIN-STRUCTURE.md", "r") as f:
    explain_prompt = f.read()

class QueryRequest(BaseModel):
    query: str

@app.post("/generate")
def generate_command(request: QueryRequest):
    messages = [{"role": "system", "content": guide_prompt}, {"role": "user", "content": request.query}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    # Automatically send inputs to whatever device the model is currently living on
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    outputs = model.generate(**inputs, max_new_tokens=150, temperature=0.1)
    response_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0].split("<|im_start|>assistant")[-1]
    
    json_start, json_end = response_text.find('{'), response_text.rfind('}') + 1
    if json_start != -1 and json_end != -1:
        try: return json.loads(response_text[json_start:json_end])
        except: pass
    return {"engine": "error", "command": "Failed to parse"}

@app.post("/explain")
def explain_command(request: QueryRequest):
    messages = [{"role": "system", "content": explain_prompt}, {"role": "user", "content": request.query}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.3)
    response_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0].split("<|im_start|>assistant")[-1]
    return {"explanation": response_text.strip()}