import streamlit as st
import requests
import os
from gliner import GLiNER
from streamlit_autorefresh import st_autorefresh

tok = os.getenv("TOK")

st_autorefresh(interval=5000, key="volter")


def Target_Identification(userinput):
    model = GLiNER.from_pretrained("Ihor/gliner-biomed-bi-small-v1.0")
    labels = ["Protein","Mutation"]
    entities = model.predict_entities(userinput, labels, threshold=0.5)

    for entity in entities:
        if entity["label"] == "Protein":
            return entity["text"]


def APP():
    tab_map = {
        0: "BIO ENGINEERING LAB @newMATTER",
    }

    tab_selection = st.pills(
        "TABS",
        options=tab_map.keys(),
        format_func=lambda option: tab_map[option],
        selection_mode="single",
    )

    def SHOWTABS():
        if tab_selection == 0:
            # Two-column split
            left_col, right_col = st.columns([0.4, 0.6])

            # CSS to make right column sticky
            st.markdown("""
                <style>
                [data-testid="column"]:nth-of-type(2) {
                    position: sticky;
                    top: 0;
                    align-self: flex-start;
                    height: 100vh;
                    overflow-y: auto;
                    background-color: #0E1117;
                    padding: 10px;
                    border-left: 1px solid rgba(255,255,255,0.1);
                }
                </style>
            """, unsafe_allow_html=True)

            with left_col:
                option_map = {
                    0: "@OriginAI  Nanobody Engineering:",
                }
                selection = st.pills(
                    "BIOLOGICS",
                    options=option_map.keys(),
                    format_func=lambda option: option_map[option],
                    selection_mode="single",
                )

                if selection == 0:
                    st.markdown(
                        "<p style='color:white;background-color:orange;font-weight:bold'> Nanobody [CANCER targeted]</p>",
                        unsafe_allow_html=True,
                    )

                    projects = []
                    projectname = None

                    def scan_for_project_availability(user_id):
                        request_url = f"https://thexforce-combat-backend.hf.space/{user_id}/projects"
                        response = requests.get(
                            request_url,
                            headers={
                                "Content-Type": "application/json",
                                "Authorization": f"Bearer {tok}",
                            },
                        )
                        response_json = response.json()
                        pros = response_json.get("projects")
                        for pro in pros:
                            if isinstance(pro, dict):
                            
                                projects.append(pro.get("project"))
                            else:
                                projects.append(pro)

                    scan_for_project_availability(st.user.email)

                    if len(projects) > 0:
                        projectname = st.selectbox("Select Project", projects)
                    else:
                        projectname = st.text_input("Enter project name:")

                    st.session_state.projectname = projectname

            with right_col:
                bio_input = st.chat_input(" Ready for Action ! ")

                @st.cache_data(ttl=10)
                def fetch_ops():
                    response = requests.get(
                        f"https://thexforce-combat-backend.hf.space/user/operations/{st.user.email}",
                        headers={
                            "Content-Type": "application/json",
                            "Authorization": f"Bearer {tok}",
                        },
                    )
                    return response.json()

                if "messages" not in st.session_state:
                    st.session_state.messages = []

                if len(st.session_state.messages) > 0:
                    for msg in st.session_state.messages:
                        with st.chat_message(msg["role"]):
                            st.markdown(msg["content"])

                if bio_input:
                    st.session_state.messages.append({"role": "user", "content": bio_input})
                    with st.chat_message("user"):
                        st.markdown(bio_input)

                    if st.session_state.projectname in [None, ""]:
                        st.markdown(":orange-badge[⚠️ Set Projectname]")
                    else:
                        identified_target = Target_Identification(bio_input)
                        st.warning(f"TARGET IDENTIFIED IS : {identified_target}")

                        payload = {
                            "uid": st.user.email,
                            "pid": st.session_state.projectname,
                            "target": identified_target or None,
                            "high_level_bio_query": bio_input,
                        }

                        response = requests.post(
                            "https://thexforce-combat-backend.hf.space/application_layer_agent",
                            json=payload,
                            headers={
                                "Content-Type": "application/json",
                                "Authorization":f"Bearer {tok}",
                            },
                        )

                        plan_response = response.json()

                        with st.chat_message("assistant"):
                            st.markdown(plan_response)

                        fetch_ops_response = fetch_ops()
                        if fetch_ops_response.get("exp") is not None:
                            for op in fetch_ops_response.get("exp"):
                                with st.chat_message("assistant"):
                                    st.markdown(op.get("operation"))
                                    st.markdown(op.get("output"))

                        st.session_state.messages.append(
                            {"role": "assistant", "content": str(plan_response)}
                        )

    if st.user.is_logged_in:
        if st.button("🚪 Logout"):
            st.logout()
            st.rerun()

        st.markdown(f"## {st.user.email}")
        SHOWTABS()
    else:
        st.info("Please log in to access the Bio Lab")
        if st.button("Log in"):
            st.login("auth0")
            st.stop()