import streamlit as st
from spacy import displacy
import spacy
import geospacy
from PIL import Image
import base64
import sys
import pandas as pd
from spacy.tokens import Span, Doc, Token
from utils import geoutil
import urllib.parse
import os
import requests
from spacy.tokens import Doc
from spacy.lang.en import English


import pydantic
print("Pydantic version:", pydantic.__version__)


API_TOKEN = os.getenv('API_TOKEN1')

BACKEND_URL = "https://SpatialWebAgent-dockerb2.hf.space/api/predict/"
def call_backend(input_text):
    try:
        headers = {
            "Authorization": f"Bearer {API_TOKEN}"
        }
        response = requests.post(
            BACKEND_URL,
            headers=headers,
            json={"data": [input_text]},
            timeout=10
        )
        if response.status_code == 200:
            return response.json()  # ✅ 保留原始 JSON 对象 (dict)
        return {"error": f"❌ Backend Error (HTTP {response.status_code})"}
    except Exception as e:
        return {"error": f"⚠️ Connection Error: {str(e)}"}




colors = {'GPE': "#43c6fc", "LOC": "#fd9720", "RSE":"#a6e22d"}
options = {"ents": ['GPE', 'LOC', "RSE"], "colors": colors}

HTML_WRAPPER = """<div style="overflow-x: auto; border: none solid #a6e22d; border-radius: 0.25rem; padding: 1rem">{}</div>"""
model = ""

gpe_selected = "GPE"
loc_selected = "LOC"
rse_selected = "RSE"

types = ""

#BASE_URL = "http://localhost:8080/"
BASE_URL = ""



def set_header():
    LOGO_IMAGE = "tetis-1.png"

    st.markdown(
        """
        <style>
        .container {
            display: flex;
        }
        .logo-text {
            font-weight:700 !important;
            font-size:50px !important;
            color: #f9a01b !important;
            padding-left: 10px !important;
        }
        .logo-img {
            float:right;
            width: 28%;
            height: 28%;
        }
        </style>
        """,
        unsafe_allow_html=True
    )
    st.markdown(
        f"""
        <div class="container">
            <img class="logo-img" src="data:image/png;base64,{base64.b64encode(open(LOGO_IMAGE, "rb").read()).decode()}">
            <p class="logo-text">GeOspaCy</p>
        </div>
        """,
        unsafe_allow_html=True
    )



def set_side_menu():

    global gpe_selected, loc_selected, rse_selected, model, types
    types =""
    params = st.query_params
    # params = st.query_params
    # print(params, 777)

    st.sidebar.markdown("## Spacy Model")
    st.sidebar.markdown("You can **select** the values of the *spacy model* from Dropdown.")
    models = ['en_core_web_sm', 'en_core_web_md', 'en_core_web_lg', 'en_core_web_trf']
    if "model" in params:
        default_ix = models.index(params["model"][0])
    else:
        default_ix = models.index('en_core_web_sm')
    model = st.sidebar.selectbox('Spacy Model',models, index=default_ix)

    st.sidebar.markdown("## Spatial Entity Labels")
    st.sidebar.markdown("**Mark** the Spatial Entities you want to extract?")
    tpes = ""
    if "type" in params:
        tpes = params['type'][0]

    if "g" in tpes:
        gpe = st.sidebar.checkbox('GPE', value = True)
    else:
        gpe = st.sidebar.checkbox('GPE')

    if "l" in tpes:
        loc = st.sidebar.checkbox('LOC', value = True)
    else:
        loc = st.sidebar.checkbox('LOC')
    if "r" in tpes:
        rse = st.sidebar.checkbox('RSE', value = True)
    else:
        rse = st.sidebar.checkbox('RSE')
    if(gpe):
        gpe_selected ="GPE"
        types+="g"

    if(loc):
        loc_selected ="LOC"
        types+="l"

    if(rse):
        rse_selected ="RSE"
        types+="r"



def set_input():
    params = st.query_params
    # params = st.query_params

    if "text" not in params:
        text = st.text_area("Input unstructured text:", "")
    else:
        text = st.text_area("Enter the text to extract {Spatial Entities}", params["text"][0])
    if(st.button("Extract")):

        return text

def set_selected_entities(doc):
    global gpe_selected, loc_selected, rse_selected, model
    ents = [ent for ent in doc.ents if ent.label_ == gpe_selected or ent.label_ == loc_selected or ent.label_ == rse_selected]

    doc.ents = ents
    return doc

def extract_spatial_entities(text):
    

    Span.set_extension("rse_id", default="", force=True)
    api_result = call_backend(text)
    print(api_result, 'dadada')
    st.markdown(api_result)
    # st.markdown(doc_element)
    doc_element = api_result["data"][0]



    nlp = English()
    nlp.add_pipe("sentencizer")
    doc = Doc(nlp.vocab).from_json(doc_element)
    doc = nlp.get_pipe("sentencizer")(doc)
    # st.markdown(type(doc))

    for ent_ext in doc_element["ents_ext"]:
        for ent in doc.ents:
            if ent.start_char == ent_ext["start"] and ent.end_char == ent_ext["end"]:
                ent._.rse_id = ent_ext["rse_id"]


    doc = set_selected_entities(doc)
    # doc.to_disk("saved_doc.spacy")
    doc.to_disk("/tmp/saved_doc.spacy")


    html = displacy.render(doc,style="ent", options = options)
    html = html.replace("\n","")
    st.write(HTML_WRAPPER.format(html),unsafe_allow_html=True)
    show_spatial_ent_table(doc, text)



    show_sentence_selector_table(doc)
    pass
def show_sentence_selector_table(doc_copy):
    st.markdown("**______________________________________________________________________________________**")
    st.markdown("**Sentence Selector for Geographic Composition**")


    sentences = list(doc_copy.sents)


    rows = []
    for idx, sent in enumerate(sentences):
        sentence_text = sent.text.strip()

        url = BASE_URL + "Tagger?mode=geocombo&text=" + urllib.parse.quote(sentence_text)
        new_row = {
            'Sr.': idx + 1,
            'sentence': sentence_text,
            'Select': f'<a target="_self" href="{url}">Select this sentence</a>'
        }
        rows.append(new_row)

    df = pd.DataFrame(rows)
    st.write(df.to_html(escape=False, index=False), unsafe_allow_html=True)



def show_spatial_ent_table(doc, text):
    global types
    if len(doc.ents) > 0:
        st.markdown("**______________________________________________________________________________________**")
        st.markdown("**Spatial Entities List**")

        df = pd.DataFrame(columns=['Sr.', 'entity', 'label', 'Map', 'GEOJson'])
        rows = [] 

        for ent in doc.ents:
            url_map = BASE_URL + "Tagger?map=true&type=" + types + "&model=" + model + "&text=" + text + "&entity=" + ent._.rse_id
            print(url_map, 'uuurrr')
            print(ent._.rse_id, 'pppp')
            url_json = BASE_URL + "Tagger?geojson=true&type=" + types + "&model=" + model + "&text=" + text + "&entity=" + ent._.rse_id

            new_row = {
                'Sr.': len(rows) + 1,
                'entity': ent.text,
                'label': ent.label_,
                'Map': f'<a target="_self" href="{url_map}">View</a>',
                'GEOJson': f'<a target="_self" href="{url_json}">View</a>'
            }

            rows.append(new_row)

        df = pd.DataFrame(rows)

        st.write(df.to_html(escape=False, index=False), unsafe_allow_html=True)

def set_header():       # tetis Geospacy LOGO
    LOGO_IMAGE = "title.jpg"

    st.markdown(
        """
        <style>
        .container {
            display: flex;
        }
        .logo-text {
            font-weight:700 !important;
            font-size:50px !important;
            color: #52aee3 !important;
            padding-left: 10px !important;
        }
        .logo-img {
            float:right;
            width: 10%;
            height: 10%;
        }
        </style>
        """,
        unsafe_allow_html=True
    )
    st.markdown(
        f"""
        <div class="container">
            <img class="logo-img" src="data:image/png;base64,{base64.b64encode(open(LOGO_IMAGE, "rb").read()).decode()}">
            <p class="logo-text">SpatialWebAgent</p>
        </div>
        """,
        unsafe_allow_html=True
    )


def set_side_menu():
    global gpe_selected, loc_selected, rse_selected, model, types
    types = ""
    params = st.query_params
    st.sidebar.markdown("## Deployment Method")
    st.sidebar.markdown("You can select the deployment method for the model.")
    deployment_options = ["API", "Local deployment"]
    use_local_model = st.sidebar.radio("Choose deployment method:", deployment_options, index=0) == "Local deployment"

    if use_local_model:
        local_model_path = st.sidebar.text_input("Enter local model path:", "")

    st.sidebar.markdown("## LLM Model")
    st.sidebar.markdown("You can **select** different  *LLM model* powered by API.")
    models = ['Llama-3-8B', 'Mistral-7B-0.3', 'Gemma-2-10B', 'GPT-4o', 'Gemini Pro', 'Deepseek-R1', 'en_core_web_sm', 'en_core_web_md', 'en_core_web_lg', 'en_core_web_trf']




    if "model" in params:
        default_ix = models.index(params["model"][0])
    else:
        default_ix = models.index('GPT-4o')




    model = st.sidebar.selectbox('LLM Model', models, index=default_ix)

    st.sidebar.markdown("## Spatial Entity Labels")

    st.sidebar.markdown("Please **Mark** the Spatial Entities you want to extract.")
    tpes = ""
    if "type" in params:
        tpes = params['type'][0]

    st.sidebar.markdown("### Absolute Spatial Entity:")
    if "g" in tpes:
        gpe = st.sidebar.checkbox('GPE', value=True)
    else:
        gpe = st.sidebar.checkbox('GPE')

    if "l" in tpes:
        loc = st.sidebar.checkbox('LOC', value=True)
    else:
        loc = st.sidebar.checkbox('LOC')

    st.sidebar.markdown("### Relative Spatial Entity:")

    if "r" in tpes:
        rse = st.sidebar.checkbox('RSE', value=True)
    else:
        rse = st.sidebar.checkbox('RSE')
    if (gpe):
        gpe_selected = "GPE"
        types += "g"

    if (loc):
        loc_selected = "LOC"
        types += "l"

    if (rse):
        rse_selected = "RSE"
        types += "r"





def main():
    global gpe_selected, loc_selected, rse_selected, model

    set_header()
    set_side_menu()


    text = set_input()

    if(text is not None):
        extract_spatial_entities(text)
    elif "text" in st.session_state:
        text = st.session_state.text
        extract_spatial_entities(text)


if __name__ == '__main__':
    main()