import json
import os
import pprint
import re

import requests
import streamlit as st
import streamlit.components.v1 as components

pp = pprint.PrettyPrinter(indent=2)
st.set_page_config(page_title="Gaia Search", layout="wide")

os.makedirs(os.path.join(os.getcwd(), ".streamlit"), exist_ok=True)
with open(os.path.join(os.getcwd(), ".streamlit/config.toml"), "w") as file:
    file.write('[theme]\nbase="light"')

LANG_MAPPING = {
    "Arabic": "ar",
    "Catalan": "ca",
    "Code": "code",
    "English": "en",
    "Spanish": "es",
    "French": "fr",
    "Indonesian": "id",
    "Indic": "indic",
    "Niger-Congo": "nigercongo",
    "Portuguese": "pt",
    "Vietnamese": "vi",
    "Chinese": "zh",
    "Detect Language": "detect_language",
    "All": "all",
}


st.sidebar.markdown(
    """
    <style>
        .aligncenter {
            text-align: center;
            font-weight: bold;
            font-size: 50px;
        }
    </style>
    <p class="aligncenter">Gaia Search 🌖🌏</p>
    <p style="text-align: center;"> A search engine for the LAION large scale image caption corpora</p>
    """,
    unsafe_allow_html=True,
)

st.sidebar.markdown(
    """
    <style>
        .aligncenter {
            text-align: center;
        }
    </style>
    <p style='text-align: center'>
    <a href="" >GitHub</a> | <a href="" >Project Report</a>
    </p>
    <p class="aligncenter">
    <a href="" target="_blank"> 
            <img src="https://colab.research.google.com/assets/colab-badge.svg"/>
        </a>
    </p>
    """,
    unsafe_allow_html=True,
)

corpus = st.sidebar.selectbox(
    "Corpus",
    (
        "LAION",
        "C4",
    ),
    index=3,
)
query = st.sidebar.text_input(label="Search query", value="")
language = st.sidebar.selectbox(
    "Language",
    (
        "Arabic",
        "Catalan",
        "Code",
        "English",
        "Spanish",
        "French",
        "Indonesian",
        "Indic",
        "Niger-Congo",
        "Portuguese",
        "Vietnamese",
        "Chinese",
        "Detect Language",
        "All",
    ),
    index=3,
)
max_results = st.sidebar.slider(
    "Maximum Number of Results",
    min_value=1,
    max_value=100,
    step=1,
    value=10,
    help="Maximum Number of Documents to return",
)
footer = """
    <style>
        .footer {
            position: fixed;
            left: 0;
            bottom: 0;
            width: 100%;
            background-color: white;
            color: black;
            text-align: center;
        }
    </style>
    <div class="footer">
        <p>
            Powered by <a href="https://huggingface.co/" >HuggingFace 🤗</a> and <a href="https://github.com/castorini/pyserini" >Pyserini 🦆</a>
        </p>
    </div>
"""
st.sidebar.markdown(footer, unsafe_allow_html=True)


def scisearch(query, corpus, language, num_results=10):
    try:
        query = query.strip()
        if query == "" or query is None:
            return

        corpus = corpus.strip()
        address = os.environ.get("address") if corpus == "LAION" else os.environ.get("address")
        post_data = {"query": query, "k": num_results}
        if language != "detect_language":
            post_data["lang"] = language

        output = requests.post(
            address, # 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"]
    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>
            """
        print(e)

    return results, highlight_terms


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 highlight_string(paragraph: str, highlight_terms: list) -> str:
    for term in highlight_terms:
        paragraph = re.sub(f"\\b{term}\\b", f"<b>{term}</b>", paragraph, flags=re.I)
    paragraph = process_pii(paragraph)
    return paragraph


def process_results(hits: list, highlight_terms: list) -> str:
    hit_list = []
    for i, hit in enumerate(hits):
        res_head = f"""
                    <div class="searchresult">
                        <h2>{i+1}. Document ID: {hit['docid']}</h2>
                        <p>Language: <string>{hit['lang']}</string>, Score: {round(hit['score'], 2)}</p>              
                    """
        for subhit in hit["meta"]["docs"]:
            res_head += f"""
                        <button onclick="load_image({subhit['_id']})">Load Image</button><br>
                        <p><img id='{subhit['_id']}' src='{subhit['URL']}'  style="width:400px;height:auto;display:none;"></p>
                        <a href='{subhit['URL']}'>{subhit['URL']}</a>
                        <p>{highlight_string(subhit['TEXT'], highlight_terms)}</p>
                        """
        res_head += f"""
                    <p>{highlight_string(hit['text'], highlight_terms)}</p>
                    </div>
                    <hr>
                    """
        hit_list.append(res_head)
    return " ".join(hit_list)


if st.sidebar.button("Search"):
    hits, highlight_terms = scisearch(query, corpus, LANG_MAPPING[language], max_results)
    html_results = process_results(hits, highlight_terms)
    rendered_results = f"""
        <div id="searchresultsarea">
            <br>
            <p id="searchresultsnumber">About {max_results} results</p>
            {html_results}
        </div>
    """
    st.markdown(
        """
        <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.0.2/dist/css/bootstrap.min.css" rel="stylesheet" 
            integrity="sha384-EVSTQN3/azprG1Anm3QDgpJLIm9Nao0Yz1ztcQTwFspd3yD65VohhpuuCOmLASjC" crossorigin="anonymous">
        """,
        unsafe_allow_html=True,
    )
    st.markdown(
        """
        <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
        """,
        unsafe_allow_html=True,
    )
    st.markdown(
        f"""
        <div class="row no-gutters mt-3 align-items-center">
            Gaia Search 🌖🌏
            <div class="col col-md-4">
                <input class="form-control border-secondary rounded-pill pr-5" type="search" value="{query}" id="example-search-input2">
            </div>
            <div class="col-auto">
                <button class="btn btn-outline-light text-dark border-0 rounded-pill ml-n5" type="button">
                    <i class="fa fa-search"></i>
                </button>
            </div>
        </div> 
        """,
        unsafe_allow_html=True,
    )
    components.html(
        """
        <style>
            #searchresultsarea {
                font-family: 'Arial';
            }
            #searchresultsnumber {
                font-size: 0.8rem;
                color: gray;
            }
            .searchresult h2 {
                font-size: 19px;
                line-height: 18px;
                font-weight: normal;
                color: rgb(7, 111, 222);
                margin-bottom: 0px;
                margin-top: 25px;
            }
            .searchresult a {
                font-size: 12px;
                line-height: 12px;
                color: green;
                margin-bottom: 0px;
            }
            .dark-mode {
                color: white;
            }
        </style>
        <script>
        function load_image(id){
            console.log(id)
            var x = document.getElementById(id);
            console.log(x)
            if (x.style.display === "none") {
                x.style.display = "block";
            } else {
                x.style.display = "none";
            }
        };
        function myFunction() {
           var element = document.body;
           element.classList.toggle("dark-mode");
        }
        </script>
        <button onclick="myFunction()">Toggle dark mode</button>
        """
        + rendered_results,
        height=800,
        scrolling=True,
    )