from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
import httpx
import asyncio
import json

app = FastAPI()

app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/prompts", StaticFiles(directory="prompts"), name="prompts")


@app.get("/")
def root():
    return FileResponse("static/index.html")


def load_config():
    with open("config.json", "r") as f:
        return json.load(f)


async def run_model(client, model_config, full_prompt):
    model_url = model_config["url"]
    model_name = model_config["name"]
    model_note = model_config["note"]
    
    augmented_prompt = f"{full_prompt}\n\nMODEL NOTE: {model_note}"
    
    try:
        async with client.stream(
            "POST",
            f"{model_url}/api/generate",
            json={
                "model": model_name,
                "prompt": augmented_prompt,
                "stream": True
            },
            timeout=120
        ) as r:
            r.raise_for_status()
            
            output = ""
            async for line in r.aiter_lines():
                if line:
                    data = json.loads(line)
                    output += data.get("response", "")

        return {
            "model": model_name,
            "output": output
        }

    except Exception as e:
        return {
            "model": model_name,
            "error": str(e)
        }


@app.post("/api/run-all")
async def run_all(payload: dict):

    prompt = payload.get("prompt", "")
    
    with open("prompts/motion-description.txt", "r") as f:
        context = f.read()

    full_prompt = f"SYSTEM:\n{context}\n\nUSER:\n{prompt}"

    config = load_config()
    models = config.get("models", [])

    async with httpx.AsyncClient() as client:
        tasks = [
            run_model(client, model, full_prompt)
            for model in models
        ]

        results = await asyncio.gather(*tasks)

    return results