import json

import streamlit as st

from pages.navigation import show_stepper
from src.fe_handler import call_predict, call_predict_all

show_stepper("Results")

st.header("Results")

input_json = st.session_state.get("input_json")

if not input_json:
    st.warning("No input provided. Go back and provide samples.")
    if st.button("Back to config", use_container_width=True):
        st.switch_page("pages/config.py")
    st.stop()

try:
    parsed = json.loads(input_json)
except (json.JSONDecodeError, ValueError) as e:
    st.error(f"Invalid input: {e}")
    if st.button("Back to config", use_container_width=True):
        st.switch_page("pages/config.py")
    st.stop()

endpoint = st.session_state.get("endpoint", "predict")

try:
    with st.spinner("Running prediction..."):
        if endpoint == "predict":
            results = call_predict(parsed)
            _single = True
        else:
            results = call_predict_all(parsed)
            _single = False

    if _single:
        for sample in results:
            st.markdown(f"**Sample {sample['id']}**")
            for entity in sample["entities"]:
                color = {"positive": "green", "neutral": "gray", "negative": "red"}.get(
                    entity["classification"], "gray"
                )
                st.markdown(f"- {entity['entity_text']}: :{color}[{entity['classification']}]")
    else:
        tabs = st.tabs(list(results.keys()))
        for tab, (mode, preds) in zip(tabs, results.items()):
            with tab:
                for sample in preds:
                    st.markdown(f"**Sample {sample['id']}**")
                    for entity in sample["entities"]:
                        color = {"positive": "green", "neutral": "gray", "negative": "red"}.get(
                            entity["classification"], "gray"
                        )
                        st.markdown(f"- {entity['entity_text']}: :{color}[{entity['classification']}]")
except Exception as e:
    st.error(f"Prediction failed: {e}")

st.write("")
if st.button("Back to config", use_container_width=True):
    st.switch_page("pages/config.py")