import os
import requests
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import FileResponse, RedirectResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

OLLAMA_API_URL = "http://localhost:11434"

app = FastAPI()

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

@app.post("/chat_api")
async def chat_endpoint(request: Request):
    body = await request.json()
    model = body.get("model", "qwen3:1.7b")
    prompt = body.get("prompt")

    if not prompt:
        raise HTTPException(status_code=400, detail="Prompt is required")

    url = f"{OLLAMA_API_URL}/api/generate"

    async def stream_response():
        try:
            response = requests.post(
                url=url,
                json={"model": model, "prompt": prompt, "stream": True},
                stream=True
            )
            for chunk in response.iter_content(chunk_size=None):
                yield chunk
        except Exception as e:
            print(f"Error during streaming: {e}")

    return StreamingResponse(stream_response(), media_type="application/x-ndjson")

@app.get("/")
async def root():
    return RedirectResponse(url="/chat")

@app.get("/chat")
async def chat_page():
    return FileResponse('chat.html')