import os
import sys
import streamlit as st
from dotenv import load_dotenv, find_dotenv
from huggingface_hub import snapshot_download  # type: ignore
import spacy
from spacy_streamlit import visualize_spans
import plotly.graph_objects as go
import json
from io import StringIO

load_dotenv(find_dotenv())
ANFORDERUNGEN = [
    "Präsenzpflicht",
    "Räumliche Flexibilität",
    "Zeitliche Flexibilität",
    "Schichtdienst",
    "Flexibiltät (SK) = Flexibilität (Soft Skill)",
]
ANGEBOTE = [
    "Planbarkeit und Verlässlichkeit",
    "Flexible und gestaltbare Dienstpläne und Schichtmodelle",
    "flexible und gestaltbare Arbeitszeiten",
    "Wahl des Arbeitszeitumfangs",
    "Job-Sharing",
    "Vereinbarkeit",
    "Unterstützung bei der Kinderbetreuung",
    "wohnortnahe Einsätze",
    "Finanzielle Unterstützung",
    "Familiäre und anlassbezogene Freistellung",
    "Unterstützung bei der Pflegeverantwortung",
]

ALL_CONCEPTS = (
    ANFORDERUNGEN
    + ANGEBOTE
    + [
        "Anforderungen an Bewerber:innen",
        "Anforderungen an Arbeitsort",
        "Zeitliche",
        "Flexibilität",
        "Angebote an Bewerber:innen",
        "Zeitsouveränität",
        "Familienfreundliche Konzepte",
    ]
)
PARENTS = [
    "Anforderungen an Arbeitsort",
    "Anforderungen an Arbeitsort",
    "Zeitliche",
    "Zeitliche",
    "Flexibilität",
    "Zeitsouveränität",
    "Zeitsouveränität",
    "Zeitsouveränität",
    "Zeitsouveränität",
    "Zeitsouveränität",
    "Familienfreundliche Konzepte",
    "Familienfreundliche Konzepte",
    "Familienfreundliche Konzepte",
    "Familienfreundliche Konzepte",
    "Familienfreundliche Konzepte",
    "Familienfreundliche Konzepte",
    "",
    "Anforderungen an Bewerber:innen",
    "Anforderungen an Bewerber:innen",
    "Anforderungen an Bewerber:innen",
    "",
    "Angebote an Bewerber:innen",
    "Angebote an Bewerber:innen",
]
fig = go.Figure(
    go.Treemap(
        labels=ALL_CONCEPTS,
        parents=PARENTS,
        textinfo="label",
    )
)
# ig.update_traces(root_color="lightblue")
fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
fig.update_layout(height=600, width=1400, template="plotly")

model_name = "de_core_news_sm"
if not spacy.util.is_package(model_name):
    # Download the model if not present
    os.system(f"python -m spacy download {model_name}")

path = snapshot_download(
    cache_dir="tmp/",
    repo_id="and-effect/family-compatibility-extractor-copy",
    revision="main",
    token=os.getenv("HF_TOKEN"),
)

sys.path.append(path)
from pipeline import PipelineWrapper  # noqa: E402


@st.cache_resource
def load_pipeline(repo_path: str) -> PipelineWrapper:
    """Load the pipeline"""
    return PipelineWrapper(path=repo_path)


@st.cache_data
def load_app_init() -> None:
    """Load the app initialisation"""
    st.title("Family Compatibility Extractor")
    st.markdown(
        """
        <style>
        .font {
        font-size:20px !important;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )

    st.markdown(
        """
        <style>
        .prediction {
            font-size:10px !important;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )
    st.markdown(
        """ <p class="font">Der Family Compatibility Extractor ist ein
        Sprachmodell, das aus den Texten von Jobanzeigen Anforderungen und
        Angebote an bzw. für Bewerbende extrahiert, die im Zusammenhang mit
        Familienkompatibilität stehen.
        Die Grafik unten zeigt die Taxonomie der extrahierten Konzepte.""",
        unsafe_allow_html=True,
    )

    # taxonomy table einfügen
    st.plotly_chart(fig)

    st.header("Untersuche das Family Compatibility Modell")
    st.markdown(
        """ <p class="font">Lade den Text einer Jobanzeige hoch oder gebe ihn
        in das Eingabefeld ein und drücke den jeweiligen Extrahieren-
        Button, um die vorhergesagten Familienkompatibilitätskonzepte und die
        dazugehörigen Erklärungen zu sehen.
        Klicke auf den Beispieltext rechts um das Eingabefeld damit zu füllen
        und das Modell zu testen.</p>""",
        unsafe_allow_html=True,
    )


st.set_page_config(layout="wide")
load_app_init()
pipeline = load_pipeline(repo_path=path)
# we need to keep button status in session state
if "clicked" not in st.session_state:
    st.session_state.clicked = {1: False, 2: False, 3: False}


def clicked(button: int) -> None:
    """updates session state variable when an example
    job title button is clicked
    Parameters
    ----------
    button : int
        The number of the button that was clicked
    """
    st.session_state.clicked[button] = True
    st.session_state.clicked[(button + 1) % 3] = False
    st.session_state.clicked[(button + 2) % 3] = False


def reset_buttons() -> None:
    """resets the session state variable for the example
    job title buttons when the user types in a job title
    """
    st.session_state.clicked = {1: False, 2: False, 3: False}


st.markdown(
    """
    <style>
    /* Style columns */
    [data-testid="column"] {
      border-radius: 15px;
         background-color: white;
         box-shadow: 0 0 10px #eee;
         border: 1px solid #ddd;
         padding: 1rem;;
    }
    /* Style containers */
    [data-testid="stVerticalBlock"] >
    [style*="flex-direction: column;"] > [data-testid="stVerticalBlock"] {
      border-radius: 15px;
         background-color: white;
         box-shadow: 0 0 10px #eee;
         border: 1px solid #ddd;
         padding: 1rem;;
    }
    </style>
    """,
    unsafe_allow_html=True,
)


st.markdown(
    """
<style>
.example {
    font-size:24px !important;
}
</style>
""",
    unsafe_allow_html=True,
)

left_column, middle_column, right_column = st.columns(3)
with right_column:
    st.markdown(
        '<p class="example">Beispiel</p>',
        unsafe_allow_html=True,
    )
    st.markdown(
        """ <p class="font">Klicke auf die Beispielausschreibung, um die
        extrahierten Konzepte anzuzeigen.</p>""",
        unsafe_allow_html=True,
    )
    st.button(
        """\n## Stellendetails zu: Mitarbeiter in der Produktion (m/w/d)\n\n### Mitarbeiter
               in der Produktion (m/w/d)\n\nHelfer/in - Elektro  \nVer\u00f6ffentlicht:
               XX.XX.2023\n\n#### Mitarbeiter in der Produktion (m/w/d)\n\nXXXX
               GmbH\n\n#### Typ:\n\nArbeit\n\n#### Arbeitszeit:\n\nVollzeit\n\n#### 
              Eintrittsdatum:\n\nab sofort\n\n#### Arbeitsort:\n\XXXX\n\n### 
              Stellenbeschreibung\n\nOrt: XXX Arbeitszeitmodell: Vollzeit 
              Vertragsart: Arbeitnehmer\u00fcberlassung mit \u00dcbernahmeoption 
              Branche: Elektro Startzeitpunkt: ab sofort Die XXXX GmbH ist Ihr
               Dienstleister im Gro\u00dfraum XXXX f\u00fcr eine faire, professionelle
               Umsetzung Ihrer neuen beruflichen Herausforderungen. Gestalten auch Sie
               Ihre berufliche Zukunft mit uns! Zum n\u00e4chstm\u00f6glichen Zeitpunkt 
              suchen wir die ideale Besetzung f\u00fcr die Position: Mitarbeiter in der 
              Produktion (m/w/d) Langfristige Besch\u00e4ftigungsoption Ihre Vorteile: 
              *All-in-One: Umfassende Beratung und unkomplizierter Bewerbungsprozess mit
               XXXX * Bezahlung: Attraktives Gehalt mit Zusatzleistungen (z.B. 
              Weihnachtsgeld) *Sicherheit: Vermittlung in einen zukunftssicheren 
              Arbeitsplatz* Work-Life-Balance: Geregelte Arbeitszeiten *Onboarding: 
              Umfangreiche und gr\u00fcndliche Einarbeitung* Abwechslung: Vielseitige
               und spannende Aufgaben *Arbeitsumgebung: Flache Hierarchien, modernes 
              Arbeitsumfeld und ein aufgeschlossenes Team* Was wir von BS  \nimmer bieten
               k\u00f6nnen: Pers\u00f6nlicher Ansprechpartner, gro\u00dfes 
              Unternehmensnetzwerk, kostenlose Arbeitskleidung sowie eine zeitnahe 
              R\u00fcckmeldung zum Bewerbungsstatus Ihre Aufgaben: *Montage von 
              Elektronik-Komponenten* Konfektionierung einzelner Komponenten 
              *Verpackungst\u00e4tigkeiten Ihr Profil:* Berufserfahrung: Idealerweise 
              erste Berufserfahrung in einer vergleichbaren Position *Hard Skills: 
              Fingerfertigkeit und technisches Grundverst\u00e4ndnis* Soft Skills: 
              Teambewusstsein und eine gute Kommunikationsf\u00e4higkeit * Arbeitsweise:
               Sorgf\u00e4ltig und strukturiert\n\n### Anforderungen an den Bewerber\n\n
              #### Berufserfahrung:\n\nMit Berufserfahrung\n\n#### F\u00e4higkeiten:\n\n*
               ##### Transport, Verkehr  \nGrundkenntnisse  \nLagerarbeit, Transport\n*
               ##### Produktion, Verarbeitung, Technik  \nGrundkenntnisse  
              \nBest\u00fccken, Maschinenf\u00fchrung, Anlagenf\u00fchrung, 
              -bedienung, Montage (Elektrotechnik)\n\n#### Pers\u00f6nliche 
              St\u00e4rken:\n\n* Selbst\u00e4ndiges Arbeiten\n* Sorgfalt/Genauigkeit\n*
               Motivation/ Leistungsbereitschaft\n* Zuverl\u00e4ssigkeit\n* 
              Lernbereitschaft\n\n#### Sprachen:\n\n* Grundkenntnisse  \nDeutsch\n\n
              #### Fahrzeug erforderlich:\n\nNein\n\n#### Reise-/Montagebereitschaft:
              \n\nNicht erforderlich\n\n### Arbeitsorte\n\n* XXXX \n\n### 
              Arbeitgeber\n\nXXX GmbH\n\n#### Firmenadresse:\n\n
              XXXX\n\n#### Branchengruppe:\n\n
              Arbeitnehmer\u00fcberlassung, Zeitarbeit\n\n#### Branche:\n\nBefristete
               \u00dcberlassung von Arbeitskr\u00e4ften\n\n### Weitere Informationen
              \n\nDas Stellenangebot wird durch den Arbeitgeber selbst verwaltet.\n\n
              Stelle im Rahmen der Arbeitnehmer\u00fcberlassung.\n\n#### Tarifvertrag:
              \n\nIGZ\n\n#### Befristung:\n\nUnbefristet\n\n#### Anzahl freier Stellen:
              \n\n1\n\n#### Arbeitszeitmodell:\n\n* Vollzeit\n\n#### Quelle des 
              Stellenangebots:\n\narbeitsagentur.de\n\n#### Referenznummer:\n
              \nXXXX\n""",
        on_click=clicked,
        args=[1],
    )
    # st.button(
    #     "Nur Vollzeit",
    #     on_click=clicked,
    #     args=[2],
    # )
    # st.button("Famillienfreundlich", on_click=clicked, args=[3])

with left_column:
    st.markdown(
        '<p class="example">Lade eine Datei mit einem Anzeigentext hoch.</p>',
        unsafe_allow_html=True,
    )
    st.markdown(
        """ <p class="font">(.txt, .json mit key "full_text")</p>""",
        unsafe_allow_html=True,
    )
    uploaded_file = st.file_uploader("Datei hochladen", type=["txt", "json"])
    if uploaded_file is not None:
        if uploaded_file.type == "text/plain":
            stringio = StringIO(uploaded_file.getvalue().decode("utf-8"))
            uploaded_jobdescription = stringio.read()
        elif uploaded_file.type == "application/json":
            file_contents = json.load(uploaded_file)
            uploaded_jobdescription = file_contents["full_text"]
        else:
            st.write("Bitte lade eine .txt oder .json Datei hoch.")

    if st.button("Extrahieren", key="extract_file_based"):
        if uploaded_jobdescription:
            query = [{"posting_id": "aaaa", "text": uploaded_jobdescription}]
            output, docs = pipeline(query)
            concept = [output[i]["concept"] for i in range(len(output))]
            if output[0]["concept"] is not None:
                st.markdown(
                    '<p class="example">Gefundene Konzepte:</p>',
                    unsafe_allow_html=True,
                )
                for c in output:
                    if c["concept"] in ANFORDERUNGEN:
                        st.write(f" - Anforderung: {c['concept']}")
                    elif c["concept"] in ANGEBOTE:
                        st.write(f" - Angebot: {c['concept']}")
                    else:
                        st.write(f" - Konzept nicht in der Taxonomie {c['concept']}")

                # add explanations
                options = {key: "#ef553b" for key in ANFORDERUNGEN} | {
                    key: "#636efa" for key in ANGEBOTE
                }
                st.markdown(
                    '<p class="example">Erklärungen:</p>',
                    unsafe_allow_html=True,
                )
                visualize_spans(
                    docs[0],
                    spans_key="ruler",
                    title="",
                    show_table=False,
                    displacy_options={"colors": options},
                )
            else:
                st.write("Es wurde kein Konzept für diese Jobanzeige extrahiert.")

        else:
            st.write("Bitte lade eine Datei hoch.")


with middle_column:
    tabs_font_css = """
    <style>
    div[class*="stTextInput"] label p {
    font-size: 2px;
    }
    </style>
    """
    st.markdown(
        '<p class="example">Gebe den Text einer Jobanzeige ein.</p>',
        unsafe_allow_html=True,
    )

    # get to know if the example job title buttons were clicked
    selected_job_text = ""
    if st.session_state.clicked[1]:
        selected_job_text = """\n## Stellendetails zu: Mitarbeiter in der Produktion (m/w/d)\n\n### Mitarbeiter
               in der Produktion (m/w/d)\n\nHelfer/in - Elektro  \nVer\u00f6ffentlicht:
               XX.XX.2023\n\n#### Mitarbeiter in der Produktion (m/w/d)\n\n XXXX
               GmbH\n\n#### Typ:\n\nArbeit\n\n#### Arbeitszeit:\n\nVollzeit\n\n#### 
              Eintrittsdatum:\n\nab sofort\n\n#### Arbeitsort:\n\XXXX\n\n### 
              Stellenbeschreibung\n\nOrt: XXXX Arbeitszeitmodell: Vollzeit 
              Vertragsart: Arbeitnehmer\u00fcberlassung mit \u00dcbernahmeoption 
              Branche: Elektro Startzeitpunkt: ab sofort Die XXXX GmbH ist Ihr
               Dienstleister im Gro\u00dfraum XXXX f\u00fcr eine faire, professionelle
               Umsetzung Ihrer neuen beruflichen Herausforderungen. Gestalten auch Sie
               Ihre berufliche Zukunft mit uns! Zum n\u00e4chstm\u00f6glichen Zeitpunkt 
              suchen wir die ideale Besetzung f\u00fcr die Position: Mitarbeiter in der 
              Produktion (m/w/d) Langfristige Besch\u00e4ftigungsoption Ihre Vorteile: 
              *All-in-One: Umfassende Beratung und unkomplizierter Bewerbungsprozess mit
               XXXX * Bezahlung: Attraktives Gehalt mit Zusatzleistungen (z.B. 
              Weihnachtsgeld) *Sicherheit: Vermittlung in einen zukunftssicheren 
              Arbeitsplatz* Work-Life-Balance: Geregelte Arbeitszeiten *Onboarding: 
              Umfangreiche und gr\u00fcndliche Einarbeitung* Abwechslung: Vielseitige
               und spannende Aufgaben *Arbeitsumgebung: Flache Hierarchien, modernes 
              Arbeitsumfeld und ein aufgeschlossenes Team* Was wir von BS  \nimmer bieten
               k\u00f6nnen: Pers\u00f6nlicher Ansprechpartner, gro\u00dfes 
              Unternehmensnetzwerk, kostenlose Arbeitskleidung sowie eine zeitnahe 
              R\u00fcckmeldung zum Bewerbungsstatus Ihre Aufgaben: *Montage von 
              Elektronik-Komponenten* Konfektionierung einzelner Komponenten 
              *Verpackungst\u00e4tigkeiten Ihr Profil:* Berufserfahrung: Idealerweise 
              erste Berufserfahrung in einer vergleichbaren Position *Hard Skills: 
              Fingerfertigkeit und technisches Grundverst\u00e4ndnis* Soft Skills: 
              Teambewusstsein und eine gute Kommunikationsf\u00e4higkeit * Arbeitsweise:
               Sorgf\u00e4ltig und strukturiert\n\n### Anforderungen an den Bewerber\n\n
              #### Berufserfahrung:\n\nMit Berufserfahrung\n\n#### F\u00e4higkeiten:\n\n*
               ##### Transport, Verkehr  \nGrundkenntnisse  \nLagerarbeit, Transport\n*
               ##### Produktion, Verarbeitung, Technik  \nGrundkenntnisse  
              \nBest\u00fccken, Maschinenf\u00fchrung, Anlagenf\u00fchrung, 
              -bedienung, Montage (Elektrotechnik)\n\n#### Pers\u00f6nliche 
              St\u00e4rken:\n\n* Selbst\u00e4ndiges Arbeiten\n* Sorgfalt/Genauigkeit\n*
               Motivation/ Leistungsbereitschaft\n* Zuverl\u00e4ssigkeit\n* 
              Lernbereitschaft\n\n#### Sprachen:\n\n* Grundkenntnisse  \nDeutsch\n\n
              #### Fahrzeug erforderlich:\n\nNein\n\n#### Reise-/Montagebereitschaft:
              \n\nNicht erforderlich\n\n### Arbeitsorte\n\n* XXXX \n\n### 
              Arbeitgeber\n\nXXX GmbH\n\n#### Firmenadresse:\n\n
              XXXX\n\n#### Branchengruppe:\n\n
              Arbeitnehmer\u00fcberlassung, Zeitarbeit\n\n#### Branche:\n\nBefristete
               \u00dcberlassung von Arbeitskr\u00e4ften\n\n### Weitere Informationen
              \n\nDas Stellenangebot wird durch den Arbeitgeber selbst verwaltet.\n\n
              Stelle im Rahmen der Arbeitnehmer\u00fcberlassung.\n\n#### Tarifvertrag:
              \n\nIGZ\n\n#### Befristung:\n\nUnbefristet\n\n#### Anzahl freier Stellen:
              \n\n1\n\n#### Arbeitszeitmodell:\n\n* Vollzeit\n\n#### Quelle des 
              Stellenangebots:\n\narbeitsagentur.de\n\n#### Referenznummer:\n
              \nXXXX\n"""

    # create a text input for the user to type in a job title
    st.text_area(
        "---",
        value=selected_job_text,
        key="typed_job_text",
        height=500,
        on_change=reset_buttons,  # noqa: E501
    )

    if selected_job_text:
        query = [{"posting_id": "aaaa", "text": selected_job_text}]
    else:
        query = [
            {"posting_id": "aaaa", "text": st.session_state.typed_job_text}
        ]  # noqa: E501
    # predicts the occupation of the job title when predict button is clicked
    if st.button("Extrahieren", key="extract_text_based"):
        output, docs = pipeline(query)
        concept = [output[i]["concept"] for i in range(len(output))]
        if output[0]["concept"] is not None:
            st.markdown(
                '<p class="example">Gefundene Konzepte:</p>',
                unsafe_allow_html=True,
            )
            for c in output:
                if c["concept"] in ANFORDERUNGEN:
                    st.write(f" - Anforderung: {c['concept']}")
                elif c["concept"] in ANGEBOTE:
                    st.write(f" - Angebot: {c['concept']}")
                else:
                    st.write(f" - Konzept nicht in der Taxonomie {c['concept']}")

            # add explanations
            options = {key: "#ef553b" for key in ANFORDERUNGEN} | {
                key: "#636efa" for key in ANGEBOTE
            }
            st.markdown(
                '<p class="example">Erklärungen:</p>',
                unsafe_allow_html=True,
            )
            visualize_spans(
                docs[0],
                spans_key="ruler",
                title="",
                show_table=False,
                displacy_options={"colors": options},
            )
        else:
            st.write("Es wurde kein Konzept für diese Jobanzeige extrahiert.")
st.markdown("#")