import streamlit as st
import json

from functions import echa_request
from functions_ui import download_pdf

st.set_page_config(
    page_title="ECHA Toxicological Data Viewer",
    page_icon="🧊",
    layout="wide",
    initial_sidebar_state="expanded"
)

if st.session_state.get('selected_cas', None) is None:
    st.warning("Nessun CAS Number selezionato. Torna alla pagina principale per effettuare una ricerca.")
    st.stop()
else:
    cas_number = st.session_state.selected_cas
    
DATA = echa_request(cas_number)

def extract_tox_info_values(data):
    """Extract DNEL values from toxicological information."""
    rows = []
    sections = data.get("toxicological_information", {}).get("sections", [])
    
    for section in sections:
        label = section.get("label", "")
        if "subsections" in section:
            for subsec in section["subsections"]:
                effect_type = subsec.get("label", "")
                if "subsections" in subsec:
                    for sub2 in subsec["subsections"]:
                        dose = sub2.get("StDose", {})
                        if isinstance(dose, dict) and dose.get("value"):
                            rows.append({
                                "Population/Route": label,
                                "Effect Type": effect_type,
                                "Exposure": sub2.get("label", ""),
                                "Assessment": sub2.get("HazardAssessment", ""),
                                "Value numerical": dose.get("value", ""),
                                "Unit": dose.get("unit", ""),
                                "Endpoint": sub2.get("MostSensitiveEndpoint", "")
                            })
    return rows


def extract_acute_values(data):
    """Extract acute toxicity values."""
    rows = []
    sections = data.get("acute_toxicity", {}).get("sections", [])
    
    for section in sections:
        if section.get("label") == "Key value for assessment":
            for subsec in section.get("subsections", []):
                if subsec.get("EffectLevelValue"):
                    rows.append({
                        "Route": subsec.get("label", "").replace("Acute toxicity: ", ""),
                        "Endpoint": subsec.get("EffectLevelUnit", ""),
                        "Value": subsec.get("EffectLevelValue", ""),
                        "Conclusion": subsec.get("EndpointConclusion", "")
                    })
    return rows


def extract_repeated_values(data):
    """Extract repeated dose toxicity values."""
    rows = []
    sections = data.get("repeated_dose_toxicity", {}).get("sections", [])
    
    for section in sections:
        if section.get("label") == "Key value for assessment":
            for subsec in section.get("subsections", []):
                study_type = subsec.get("label", "")
                for sub2 in subsec.get("subsections", []):
                    if sub2.get("EffectLevelValue"):
                        rows.append({
                            "Study Type": study_type,
                            "Route": sub2.get("label", ""),
                            "Endpoint": sub2.get("EffectLevelUnit", ""),
                            "Value": sub2.get("EffectLevelValue", ""),
                            "Species": sub2.get("Species", "-")
                        })
    return rows


# App
st.title("Toxicological Data Viewer")

# Substance info
substance = DATA["substance"]
dossier = DATA["dossier_info"]

st.subheader(f"{substance['rmlName']}")

with st.container(border=True):
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.caption("CAS Number")
        st.write(substance['rmlCas'])
    with col2:
        st.caption("EC Number")
        st.write(substance['rmlEc'])
    with col3:
        st.caption("Status")
        st.write(dossier['registrationStatus'])
    with col4:
        st.caption("Last Updated")
        st.write(dossier['lastUpdatedDate'])

# Tabs
tab1, tab2, tab3 = st.tabs(["Toxicological Information", "Acute Toxicity", "Repeated Dose Toxicity"])

with tab1:
    st.subheader("Derived No Effect Levels (DNEL)")
    tox_rows = extract_tox_info_values(DATA)
    if tox_rows:
        st.table(tox_rows)
    else:
        st.info("No DNEL values found.")
        
    download_pdf(
        casNo=substance['rmlCas'],
        origin='echa_tox_info',
        link=DATA['index']['toxicological_information_link']
    )

with tab2:
    st.subheader("Acute Toxicity Values")
    acute_rows = extract_acute_values(DATA)
    if acute_rows:
        st.table(acute_rows)
    else:
        st.info("No acute toxicity values found.")
        
    download_pdf(
        casNo=substance['rmlCas'],
        origin='echa_acute_tox',
        link=DATA['index']['acute_toxicity_link']
    )

with tab3:
    st.subheader("Repeated Dose Toxicity Values")
    repeated_rows = extract_repeated_values(DATA)
    if repeated_rows:
        st.table(repeated_rows)
    else:
        st.info("No repeated dose toxicity values found.")
    
    download_pdf(
        casNo=substance['rmlCas'],
        origin='echa_repeated_tox',
        link=DATA['index']['repeated_dose_toxicity_link']
    )
    


# Key Information sections
st.divider()
st.subheader("Key Information")

acute_key_info = None
repeated_key_info = None

for section in DATA.get("acute_toxicity", {}).get("sections", []):
    if section.get("label") == "Description of key information":
        acute_key_info = section.get("KeyInformation")
        break

for section in DATA.get("repeated_dose_toxicity", {}).get("sections", []):
    if section.get("label") == "Description of key information":
        repeated_key_info = section.get("KeyInformation")
        break

if acute_key_info:
    with st.expander("Acute Toxicity - Key Information"):
        st.write(acute_key_info)

if repeated_key_info:
    with st.expander("Repeated Dose Toxicity - Key Information"):
        st.write(repeated_key_info)