from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from typing import Optional

from prompts import get_summary_prompt
from _llm import LLM
from _arango import ArangoDB

app = FastAPI()


class DocumentData(BaseModel):
    text: str
    arango_db_name: str
    arango_id: str
    is_sci: Optional[bool] = False


@app.post("/summarise_document")
async def summarize_document(doc_data: DocumentData, background_tasks: BackgroundTasks):
    background_tasks.add_task(summarise_document_task, doc_data.dict())
    return {"message": "Document summarization has started."}


def summarise_document_task(doc_data: dict):
    text = doc_data.get("text")
    is_sci = doc_data.get("is_sci", False)

    system_message = "You are summarising scientific articles. It is very important that you keep to what is written and do not add any of your own opinions or interpretations. Always answer in English."
    llm = LLM(system_message=system_message)

    summary = llm.generate(query=get_summary_prompt(text, is_sci))

    summary_doc = {
        "text_sum": summary,
        "meta": {
            "model": llm.model,
            "system_message": system_message,
            "temperature": llm.options["temperature"],
        },
    }

    arango = ArangoDB(db_name=doc_data.get("arango_db_name"))
    arango.db.update_document(
        {"summary": summary_doc, "_id": doc_data.get("arango_id")},
        silent=True,
        check_rev=False,
    )


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