from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import mlx_lm
import os
import time
from typing import List, Optional
from pydantic import BaseModel

# --- Fast, Apple Silicon Native Inference ---
# Uses mlx-lm library which is highly optimized for 16GB Mac RAM.

app = FastAPI(title="Ankahi Mac Server", version="1.0.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

MODEL_PATH = "ankahi_mlx_4bit" # The 4-bit optimized model folder

class PredictRequest(BaseModel):
    pictogram_sequence: List[str]
    persona_id: str
    n_alternatives: int = 3

model = None
tokenizer = None

@app.on_event("startup")
async def load_optimized_model():
    global model, tokenizer
    print(f"Loading 4-bit optimized model from {MODEL_PATH}...")
    try:
        model, tokenizer = mlx_lm.load(MODEL_PATH)
        print("Optimized model loaded on Apple Silicon!")
    except Exception as e:
        print(f"Error: Could not load model. Did you run convert_to_mlx.py first? {e}")

@app.post("/predict")
async def predict(request: PredictRequest):
    if not model:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    start_time = time.time()
    prompt = f"Pictogram sequence: [{', '.join(request.pictogram_sequence)}]. Predict {request.n_alternatives} sentences."
    
    # Fast Apple Silicon Inference
    response = mlx_lm.generate(model, tokenizer, prompt=prompt, max_tokens=150)
    
    # Simple list parsing for demo (real version uses regex)
    alternatives = [line.strip() for line in response.split('\n') if line.strip()][:request.n_alternatives]
    
    return {
        "alternatives": alternatives or [response],
        "primary": alternatives[0] if alternatives else response,
        "inference_time_ms": (time.time() - start_time) * 1000,
        "persona_id": request.persona_id
    }

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