from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
from llm_manager import LLMManager
import json
import time
import asyncio
from typing import Optional

app = FastAPI()
llm = None

@app.on_event("startup")
async def startup_event():
    global llm
    llm = LLMManager()

@app.get("/")
async def root():
    return {"status": "running", "model": llm.model_file if llm else "loading"}

@app.post("/api/generate")
async def generate(request: Request):
    data = await request.json()
    prompt = data.get("prompt")
    stream = data.get("stream", True)
    model_name = data.get("model", "qwen3")

    if not prompt:
        return JSONResponse({"error": "Prompt is required"}, status_code=400)

    def stream_response():
        response = llm.generate(prompt, stream=True)
        for chunk in response:
            yield json.dumps({
                "model": model_name,
                "created_at": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
                "response": chunk["choices"][0]["text"],
                "done": False
            }) + "\n"
        
        yield json.dumps({
            "model": model_name,
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
            "done": True,
            "context": [], # Placeholder
            "total_duration": 0,
            "load_duration": 0,
            "prompt_eval_count": 0,
            "prompt_eval_duration": 0,
            "eval_count": 0,
            "eval_duration": 0
        }) + "\n"

    if stream:
        return StreamingResponse(stream_response(), media_type="application/x-ndjson")
    else:
        response = llm.generate(prompt, stream=False)
        return {
            "model": model_name,
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
            "response": response["choices"][0]["text"],
            "done": True,
            "context": [],
            "total_duration": 0,
            "load_duration": 0,
            "prompt_eval_count": 0,
            "prompt_eval_duration": 0,
            "eval_count": 0,
            "eval_duration": 0
        }

@app.post("/api/chat")
async def chat(request: Request):
    data = await request.json()
    messages = data.get("messages", [])
    stream = data.get("stream", True)
    model_name = data.get("model", "qwen3")

    def stream_chat():
        response = llm.chat_completion(messages, stream=True)
        for chunk in response:
            if "choices" in chunk and len(chunk["choices"]) > 0:
                delta = chunk["choices"][0].get("delta", {})
                content = delta.get("content", "")
                yield json.dumps({
                    "model": model_name,
                    "created_at": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
                    "message": {"role": "assistant", "content": content},
                    "done": False
                }) + "\n"
        
        yield json.dumps({
            "model": model_name,
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
            "done": True
        }) + "\n"

    if stream:
        return StreamingResponse(stream_chat(), media_type="application/x-ndjson")
    else:
        response = llm.chat_completion(messages, stream=False)
        return {
            "model": model_name,
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
            "message": response["choices"][0]["message"],
            "done": True
        }

@app.get("/api/tags")
async def tags():
    return {
        "models": [
            {
                "name": "qwen3",
                "modified_at": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
                "size": 0, # TBD
                "digest": "qwen3-digest",
                "details": {
                    "format": "gguf",
                    "family": "qwen",
                    "families": ["qwen"],
                    "parameter_size": "14B",
                    "quantization_level": "Q4_K_M"
                }
            }
        ]
    }

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