import os

from fastapi import FastAPI, WebSocket
from epub2txt import epub2txt

from fastllm_pytools import llm

model = llm.model(os.getenv("checkpoint_path"))
prompt = os.getenv("prompt")

app = FastAPI()


@app.websocket("/ws")
async def read_root(websocket: WebSocket):
    await websocket.accept()
    f_name = await websocket.receive_text()
    ch_list = epub2txt(f_name, outputlist=True)
    chapter_titles = epub2txt.content_titles
    title = epub2txt.title

    idx = 0
    sm_list = []
    for text in ch_list[2:]:
        idx += 1
        docs = []
        for i in range(0, len(text)//2000+1, 2000):
            t = text[i:i+2048]
            if len(t) > 0:
                docs.append(model.response(prompt+t))
                await websocket.send_text(f"chsum: {docs[-1]}")
        hist = docs[0]
        for doc in docs[1:]:
            hist = model.response(prompt+"\n"+hist+"\n"+doc)
            await websocket.send_text(f"draft_sum: {hist}")
        sm_list.append(hist)
        mdobj_str = f"# {title}\n\n{hist}\n\n\n"
        for ct, sm in zip(chapter_titles[2:], sm_list):
            mdobj_str += f"## {ct}\n\n{sm}\n\n\n"
        await websocket.send_text(f"output: {mdobj_str}")

# uvicorn api:app --reload