from fastapi.responses import JSONResponse
from fastapi import APIRouter, status
from openai import OpenAI
import os

API_KEY = os.environ.get("openai_api_key")
ASSISTANT_ID = os.environ.get("openai_assistant_id")
router = APIRouter()
client = OpenAI(api_key=API_KEY)


@router.get("/openai")
async def call_openai(prompt: str):
    thread = client.beta.threads.create(
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt,
                    }
                ],
            }
        ]
    )
    run = client.beta.threads.runs.create_and_poll(
        thread_id=thread.id,
        assistant_id=ASSISTANT_ID,
    )
    start_time = run.started_at
    while True:
        messages = list(
            client.beta.threads.messages.list(thread_id=thread.id, run_id=run.id)
        )
        if messages:
            print("received {} messages".format(len(messages)))
            break
        end_time = run.completed_at
        if end_time - start_time > 60:
            break
    if not messages:
        return JSONResponse(
            status_code=status.HTTP_408_REQUEST_TIMEOUT,
            content={
                "code": run.code,
                "message": run.message,
            },
        )
    message_content = messages[0].content[0].text
    annotations = message_content.annotations
    citations = []
    for index, annotation in enumerate(annotations):
        message_content.value = message_content.value.replace(
            annotation.text, f"[{index}]"
        )
        file_citation = getattr(annotation, "file_citation", None)
        if file_citation:
            cited_file = client.files.retrieve(file_citation.file_id)
            citations.append(f"[{index}] {cited_file.filename}")
    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content={
            "assistant_id": ASSISTANT_ID,
            "thread_id": thread.id,
            "run_id": run.id,
            "answer": message_content.value,
            "citations": citations,
        },
    )