from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware

api = FastAPI()

# Enable traffic parsing variables
api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@api.post("/v1/chat/completions")
async def handle_opencode_request(request: Request):
    try:
        json_data = await request.json()
        messages = json_data.get("messages", [])
        user_prompt = messages[-1]["content"] if messages else ""
        chosen_model = json_data.get("model", "bonsai-27b")
    except Exception:
        return JSONResponse({"error": "Invalid JSON format"}, status_code=400)
    
    # Imports functions dynamically from your live main script layer
    from app import generate_27b, generate_moe_cpu
    
    if "liquid" in chosen_model or "cpu" in chosen_model or "backup" in chosen_model:
        model_reply = generate_moe_cpu(user_prompt)
    else:
        model_reply = generate_27b(user_prompt)
        
    return JSONResponse({
        "id": "hf-opencode-session",
        "object": "chat.completion",
        "model": chosen_model,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": model_reply},
            "finish_reason": "stop"
        }]
    })

# Add an automatic redirect map so the root url resolves correctly
@api.get("/")
def read_root():
    return {"status": "OpenAI Companion Server Active on Port 8000"}