import os
import subprocess
import gradio as gr
from epub2txt import epub2txt
from websocket import create_connection

if not os.path.exists(os.getenv("checkpoint_path")):
    os.system("git clone --recurse-submodules https://github.com/ztxz16/fastllm.git")
    os.system("cd fastllm; mkdir build; cd build; cmake ..; make -j; cd tools; python setup.py install --user --prefix=")
    os.system("wget https://huggingface.co/huangyuyang/chatglm2-6b-int4.flm/resolve/main/chatglm2-6b-int4.flm")
    subprocess.Popen(["uvicorn", "api:app"])

class GUI:
    def __init__(self, *args, **kwargs):
        with gr.Blocks() as demo:
            with gr.Row():
                gr.Markdown(scale=2).attach_load_event(self.hello, None)
                gr.LoginButton()
                gr.LogoutButton()
            out = gr.Markdown()
            inp = gr.File(file_types=['.epub'])
            inp.change(self.process, inp, out)
            self.ws = None
            self.out = []
        demo.queue(concurrency_count=2).launch()

    def process(self, file, profile: gr.OAuthProfile | None):
        if profile is None:
            return gr.update(value='Login to access the tool.')

        h = ''
        chapter_titles = epub2txt.content_titles
        title = epub2txt.title
        if self.ws is None:
            self.ws = create_connection(f"ws://localhost:8000/ws")
            self.ws.send(file.name)
        res = ''
        while 'output: ' not in res:
            res = self.ws.recv()
            if 'chsum: ' in res:
                self.out.append(res.replace("chsum: ", ""))
            elif 'draft_sum: ' in res:
                h = res[11:]
            elif 'output: ' in res:
                self.ws.close()
                self.ws = None
                self.out = []
                yield gr.update(value=res)
            yield gr.update(
                value=f"# {title}\n\n" +h+ "\n\n".join(
                    [f"## {ct}\n\n{c}" for ct, c in zip(chapter_titles, self.out)]))

    def hello(self, profile: gr.OAuthProfile | None):
        if profile is None:
            return '# ePub summarization tool\n\nLogin to access the tool.'
        return f"# ePub summarization tool\n\nWelcome {profile.name}!!"

GUI()