import subprocess

subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
               shell=True)

import gradio as gr
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)
import docx
import PyPDF2
import spaces


def convert_to_txt(file):
    doc_type = file.split(".")[-1].strip()
    if doc_type in ["txt", "md", "py"]:
        data = [file.read().decode("utf-8")]
    elif doc_type in ["pdf"]:
        pdf_reader = PyPDF2.PdfReader(file)
        data = [
            pdf_reader.pages[i].extract_text() for i in range(len(pdf_reader.pages))
        ]
    elif doc_type in ["docx"]:
        doc = docx.Document(file)
        data = [p.text for p in doc.paragraphs]
    else:
        raise gr.Error(f"ERROR: unsupported document type: {doc_type}")
    text = "\n\n".join(data)
    return text


model_name = "THUDM/LongCite-glm4-9b"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device="cuda",
    attn_implementation="flash_attention_2",
)

html_styles = """<style>
    .reference {
        color: blue;
        text-decoration: underline;
    }
    .highlight {
        background-color: yellow;
    }
    .label {
        font-family: sans-serif;
        font-size: 16px;
        font-weight: bold;
    }
    .Bold {
        font-weight: bold;
    }
    .statement {
        background-color: lightgrey;
    }
</style>\n"""


def process_text(text):
    special_char = {
        "&": "&amp;",
        "'": "&apos;",
        '"': "&quot;",
        "<": "&lt;",
        ">": "&gt;",
        "\n": "<br>",
    }
    for x, y in special_char.items():
        text = text.replace(x, y)
    return text


def convert_to_html(statements, clicked=-1):
    html = html_styles + '<br><span class="label">Answer:</span><br>\n'
    all_cite_html = []
    clicked_cite_html = None
    cite_num2idx = {}
    idx = 0
    for i, js in enumerate(statements):
        statement, citations = process_text(js["statement"]), js["citation"]
        if clicked == i:
            html += f"""<span class="statement">{statement}</span>"""
        else:
            html += f"<span>{statement}</span>"
        if citations:
            cite_html = []
            idxs = []
            for c in citations:
                idx += 1
                idxs.append(str(idx))
                cite = (
                    "[Sentence: {}-{}\t|\tChar: {}-{}]<br>\n<span {}>{}</span>".format(
                        c["start_sentence_idx"],
                        c["end_sentence_idx"],
                        c["start_char_idx"],
                        c["end_char_idx"],
                        'class="highlight"' if clicked == i else "",
                        process_text(c["cite"].strip()),
                    )
                )
                cite_html.append(
                    f"""<span><span class="Bold">Snippet [{idx}]:</span><br>{cite}</span>"""
                )
            all_cite_html.extend(cite_html)
            cite_num = "[{}]".format(",".join(idxs))
            cite_num2idx[cite_num] = i
            cite_num_html = """ <span class="reference" style="color: blue" id={}>{}</span>""".format(
                i, cite_num
            )
            html += cite_num_html
        html += "\n"
        if clicked == i:
            clicked_cite_html = (
                    html_styles
                    + """<br><span class="label">Citations of current statement:</span><br><div style="overflow-y: auto; padding: 20px; border: 0px dashed black; border-radius: 6px; background-color: #EFF2F6;">{}</div>""".format(
                "<br><br>\n".join(cite_html)
            )
            )
    all_cite_html = (
            html_styles
            + """<br><span class="label">All citations:</span><br>\n<div style="overflow-y: auto; padding: 20px; border: 0px dashed black; border-radius: 6px; background-color: #EFF2F6;">{}</div>""".format(
        "<br><br>\n".join(all_cite_html).replace(
            '<span class="highlight">', "<span>"
        )
        if len(all_cite_html)
        else "No citation in the answer"
    )
    )
    return html, all_cite_html, clicked_cite_html, cite_num2idx


def render_context(file):
    if hasattr(file, "name"):
        context = convert_to_txt(file.name)
        return gr.Textbox(context, visible=True)
    else:
        raise gr.Error(f"ERROR: no uploaded document")


@spaces.GPU(duration=120)
def infer(context, query):
    return model.query_longcite(
        context=context,
        query=query,
        tokenizer=tokenizer,
        max_input_length=128000,
        max_new_tokens=1024,
    )

def run_llm(context, query):
    if not context:
        raise gr.Error("Error: no uploaded document")
    if not query:
        raise gr.Error("Error: no query")
    result = infer(context=context, query=query)
    all_statements = result["all_statements"]
    answer_html, all_cite_html, clicked_cite_html, cite_num2idx_dict = convert_to_html(
        all_statements
    )
    cite_nums = list(cite_num2idx_dict.keys())
    return {
        statements: gr.JSON(all_statements),
        answer: gr.HTML(answer_html, visible=True),
        all_citations: gr.HTML(all_cite_html, visible=True),
        cite_num2idx: gr.JSON(cite_num2idx_dict),
        citation_choices: gr.Radio(cite_nums, visible=len(cite_nums) > 0),
        clicked_citations: gr.HTML(visible=False),
    }


def chose_citation(statements, cite_num2idx, clicked_cite_num):
    clicked = cite_num2idx[clicked_cite_num]
    answer_html, _, clicked_cite_html, _ = convert_to_html(statements, clicked=clicked)
    return {
        answer: gr.HTML(answer_html, visible=True),
        clicked_citations: gr.HTML(clicked_cite_html, visible=True),
    }


with gr.Blocks() as demo:
    gr.Markdown(
        """
        <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
            LongCite-glm4-9b Huggingface Space🤗
        </div>
        <div style="text-align: center;">
            <a href="https://huggingface.co/THUDM/LongCite-glm4-9b">🤗 Model Hub</a> |
            <a href="https://github.com/THUDM/LongCite">🌐 Github</a> |
            <a href="https://arxiv.org/abs/2409.02897">📜 arxiv </a>
        </div>
        <br>
        <div style="text-align: center; font-size: 15px; font-weight: bold; margin-bottom: 20px; line-height: 1.5;">
        If you plan to use it long-term, please consider deploying the model or forking this space yourself.
        </div>  
        """
    )

    with gr.Row():
        with gr.Column(scale=4):
            file = gr.File(
                label="Upload a document (supported type: pdf, docx, txt, md, py)"
            )
            query = gr.Textbox(label="Question")
            submit_btn = gr.Button("Submit")

        with gr.Column(scale=4):
            context = gr.Textbox(
                label="Document content",
                autoscroll=False,
                placeholder="No uploaded document.",
                max_lines=10,
                visible=False,
            )

            file.upload(render_context, [file], [context])

    with gr.Row():
        with gr.Column(scale=4):
            statements = gr.JSON(label="statements", visible=False)
            answer = gr.HTML(label="Answer", visible=True)
            cite_num2idx = gr.JSON(label="cite_num2idx", visible=False)
            citation_choices = gr.Radio(
                label="Chose citations for details", visible=False, interactive=True
            )

        with gr.Column(scale=4):
            clicked_citations = gr.HTML(
                label="Citations of the chosen statement", visible=False
            )
            all_citations = gr.HTML(label="All citations", visible=False)

    submit_btn.click(
        run_llm,
        [context, query],
        [
            statements,
            answer,
            all_citations,
            cite_num2idx,
            citation_choices,
            clicked_citations,
        ],
    )
    citation_choices.change(
        chose_citation,
        [statements, cite_num2idx, citation_choices],
        [answer, clicked_citations],
    )

demo.queue()
demo.launch()