import json
import os
import gradio as gr
import requests
from huggingface_hub import HfApi

hf_api = HfApi()
roots_datasets = {
    dset.id.split("/")[-1]: dset
    for dset in hf_api.list_datasets(
        author="bigscience-data", use_auth_token=os.environ.get("bigscience_data_token")
    )
}


def get_docid_html(docid):
    data_org, dataset, docid = docid.split("/")
    metadata = roots_datasets[dataset]
    if metadata.private:
        docid_html = (
            f"<a "
            f'class="underline-on-hover"'
            f'title="This dataset is private. See the introductory text for more information"'
            f'style="color:#AA4A44;"'
            f'href="https://huggingface.co/datasets/bigscience-data/{dataset}"'
            f'target="_blank"><b>🔒{dataset}</b></a><span style="color: #7978FF;">/{docid}</span>'
        )
    else:
        docid_html = (
            f"<a "
            f'class="underline-on-hover"'
            f'title="This dataset is licensed {metadata.tags[0].split(":")[-1]}"'
            f'style="color:#2D31FA;"'
            f'href="https://huggingface.co/datasets/bigscience-data/{dataset}"'
            f'target="_blank"><b>{dataset}</b></a><span style="color: #7978FF;">/{docid}</span>'
        )
    return docid_html


PII_TAGS = {"KEY", "EMAIL", "USER", "IP_ADDRESS", "ID", "IPv4", "IPv6"}
PII_PREFIX = "PI:"


def process_pii(text):
    for tag in PII_TAGS:
        text = text.replace(
            PII_PREFIX + tag,
            """<b><mark style="background: Fuchsia; color: Lime;">REDACTED {}</mark></b>""".format(
                tag
            ),
        )
    return text


def process_results(results, highlight_terms):
    if len(results) == 0:
        return """<br><p style='font-family: Arial; color:Silver; text-align: center;'>
                No results retrieved.</p><br><hr>"""

    results_html = ""
    for result in results:
        tokens = result["text"].split()
        tokens_html = []
        for token in tokens:
            if token in highlight_terms:
                tokens_html.append("<b>{}</b>".format(token))
            else:
                tokens_html.append(token)
        tokens_html = " ".join(tokens_html)
        tokens_html = process_pii(tokens_html)
        meta_html = (
            """
                <p class='underline-on-hover' style='font-size:12px; font-family: Arial; color:#585858; text-align: left;'>
                <a href='{}' target='_blank'>{}</a></p>""".format(
                result["meta"]["url"], result["meta"]["url"]
            )
            if "meta" in result
            and result["meta"] is not None
            and "url" in result["meta"]
            else ""
        )
        docid_html = get_docid_html(result["docid"])
        results_html += """{}
            <p style='font-size:14px; font-family: Arial; color:#7978FF; text-align: left;'>Document ID: {}</p>
            <p style='font-size:12px; font-family: Arial; color:MediumAquaMarine'>Language: {}</p>
            <p style='font-family: Arial;'>{}</p>
            <br>
        """.format(
            meta_html, docid_html, result["lang"], tokens_html
        )
    return results_html + "<hr>"


def scisearch(query, language, num_results=10, exact_search=False):
    try:
        query = " ".join(query.split())
        if query == "" or query is None:
            return ""

        post_data = {"query": query, "k": num_results, "exact_search": exact_search}
        if language != "detect_language":
            post_data["lang"] = language

        output = requests.post(
            os.environ.get("address"),
            headers={"Content-type": "application/json"},
            data=json.dumps(post_data),
            timeout=60,
        )

        payload = json.loads(output.text)

        if "err" in payload:
            if payload["err"]["type"] == "unsupported_lang":
                detected_lang = payload["err"]["meta"]["detected_lang"]
                return f"""
                    <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'>
                    Detected language <b>{detected_lang}</b> is not supported.<br>
                    Please choose a language from the dropdown or type another query.
                    </p><br><hr><br>"""

        results = payload["results"]
        highlight_terms = payload["highlight_terms"]

        if language == "detect_language":
            return (
                (
                    f"""<p style='font-family: Arial; color:MediumAquaMarine; text-align: center; line-height: 3em'>
                Detected language: <b>{results[0]["lang"]}</b></p><br><hr><br>"""
                    if len(results) > 0 and language == "detect_language"
                    else ""
                )
                + process_results(results, highlight_terms)
            )

        if language == "all":
            results_html = ""
            for lang, results_for_lang in results.items():
                if len(results_for_lang) == 0:
                    results_html += f"""<p style='font-family: Arial; color:Silver; text-align: left; line-height: 3em'>
                            No results for language: <b>{lang}</b><hr></p>"""
                    continue

                collapsible_results = f"""
                    <details>
                        <summary style='font-family: Arial; color:MediumAquaMarine; text-align: left; line-height: 3em'>
                            Results for language: <b>{lang}</b><hr>
                        </summary>
                        {process_results(results_for_lang, highlight_terms)}
                    </details>"""
                results_html += collapsible_results
            return results_html

        return process_results(results, highlight_terms)

    except Exception as e:
        results_html = f"""
                <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'>
                Raised {type(e).__name__}</p>
                <p style='font-size:14px; font-family: Arial; '>
                Check if a relevant discussion already exists in the Community tab. If not, please open a discussion.
                </p>
            """

    return results_html


def perform_exact_search(query, num_results=10):
    try:
        print("perform_exact_search")
        query = " ".join(query.split())
        if query == "" or query is None:
            return ""

        post_data = {"query": query, "k": num_results, "exact_search": True}

        print("post_data", post_data)

        output = requests.post(
            "http://34.105.160.81:8080",
            headers={"Content-type": "application/json"},
            data=json.dumps(post_data),
            timeout=60,
        )

        payload = json.loads(output.text)
        results = payload["results"]

        results_html = ""
        for result in results:
            print(result)
            result_html = """<br><hr><br>"""
            query_start = result.find(query)
            query_end = query_start + len(query)
            result_html += result[0:query_start]
            result_html += "<b>{}</b>".format(result[query_start:query_end])
            result_html += result[query_end:]
            results_html += result_html
        return results_html + "<hr>"

    except Exception as e:
        results_html = f"""
                <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'>
                Raised {type(e).__name__}</p>
                <p style='font-size:14px; font-family: Arial; '>
                Check if a relevant discussion already exists in the Community tab. If not, please open a discussion.
                </p>
            """


def flag(query, language, num_results, issue_description):
    try:
        post_data = {
            "query": query,
            "k": num_results,
            "flag": True,
            "description": issue_description,
        }
        if language != "detect_language":
            post_data["lang"] = language

        output = requests.post(
            os.environ.get("address"),
            headers={"Content-type": "application/json"},
            data=json.dumps(post_data),
            timeout=120,
        )

        results = json.loads(output.text)
    except:
        print("Error flagging")
    return ""


description = """# <p style="text-align: center;"> 🌸 🔎 ROOTS search tool 🔍 🌸 </p>
The ROOTS corpus was developed during the [BigScience workshop](https://bigscience.huggingface.co/) for the purpose
of training the Multilingual Large Language Model [BLOOM](https://huggingface.co/bigscience/bloom). This tool allows
you to search through the ROOTS corpus. We serve a BM25 index for each language or group of languages included in
ROOTS. You can read more about the details of the tool design
[here](https://huggingface.co/spaces/bigscience-data/scisearch/blob/main/roots_search_tool_specs.pdf). For more
information and instructions on how to access the full corpus check [this form](https://forms.gle/qyYswbEL5kA23Wu99)."""


if __name__ == "__main__":
    demo = gr.Blocks(
        css=".underline-on-hover:hover { text-decoration: underline; } .flagging { font-size:12px; color:Silver; }"
    )

    with demo:
        with gr.Row():
            gr.Markdown(value=description)
        with gr.Row():
            query = gr.Textbox(
                lines=1,
                max_lines=1,
                placeholder="Type your query here...",
                label="Query",
            )
        with gr.Row():
            lang = gr.Dropdown(
                choices=[
                    "ar",
                    "ca",
                    "code",
                    "en",
                    "es",
                    "eu",
                    "fr",
                    "id",
                    "indic",
                    "nigercongo",
                    "pt",
                    "vi",
                    "zh",
                    "detect_language",
                    "all",
                ],
                value="en",
                label="Language",
            )
        with gr.Row():
            k = gr.Slider(1, 100, value=10, step=1, label="Max Results")
        with gr.Row():
            with gr.Column(scale=1):
                exact_search = gr.Checkbox(
                    value=False, label="Exact Search", variant="compact"
                )
            with gr.Column(scale=4):
                submit_btn = gr.Button("Submit")
        with gr.Row():
            results = gr.HTML(label="Results")
        flag_description = """
            <p class='flagging'>
            If you choose to flag your search, we will save the query, language and the number of results you requested.
            Please consider adding any additional context in the box on the right.</p>"""
        with gr.Column(visible=False) as flagging_form:
            flag_txt = gr.Textbox(
                lines=1,
                placeholder="Type here...",
                label="""If you choose to flag your search, we will save the query, language and the number of results
                    you requested. Please consider adding relevant additional context below:""",
            )
            flag_btn = gr.Button("Flag Results")
            flag_btn.click(flag, inputs=[query, lang, k, flag_txt], outputs=[flag_txt])

        def submit(query, lang, k, exact_search):
            print("submitting", query, lang, k, exact_search)
            query = query.strip()
            if query is None or query == "":
                return "", ""

            if exact_search:
                return {
                    results: perform_exact_search(query, k),
                    flagging_form: gr.update(visible=True),
                }
            return {
                results: scisearch(query, lang, k, exact_search),
                flagging_form: gr.update(visible=True),
            }

        query.submit(
            fn=submit,
            inputs=[query, lang, k, exact_search],
            outputs=[results, flagging_form],
        )
        submit_btn.click(
            submit,
            inputs=[query, lang, k, exact_search],
            outputs=[results, flagging_form],
        )

    demo.launch(enable_queue=False, debug=True)