import streamlit as st

from graph_pipeline import textdoctor_app, TextDoctorState


st.set_page_config(
    page_title="TextDoctor – Multi-Agent Text Polisher",
    page_icon="🩺",
    layout="wide"
)

st.title("🩺 TextDoctor – Multi-Agent Text Improvement")
st.write(
    "Give me your **informal / messy text**, and I’ll run it through "
    "grammar, style, and clarity agents, then a reviewer agent will pick "
    "the best professional version with explanations."
)

# --- Input area ---
default_example = (
    "bro i need 2 more days for this, hope thats okay, "
    "i had some issues so couldnt finish on time"
)

user_text = st.text_area(
    "Enter your text here:",
    value=default_example,
    height=150
)

col1, _ = st.columns([1, 3])

with col1:
    run_button = st.button("Run TextDoctor 🚀")

if run_button and user_text.strip():
    with st.spinner("Running multi-agent pipeline..."):
        try:
            # Prepare initial state for LangGraph app
            state: TextDoctorState = {"input_text": user_text}
            final_state = textdoctor_app.invoke(state)
        except Exception as e:
            st.error(f"Something went wrong while running the pipeline: {e}")
            st.stop()

    review = final_state["review_result"]
    final_text = review["final_text"]

    st.subheader("✅ Final Professional Version")
    st.success(final_text)

    with st.expander("🧠 Why this version? (Reviewer Agent Explanation)", expanded=False):
        st.write(review["decision_explanation"])

    # --- Show intermediate steps ---
    st.subheader("🔍 Agent-wise Transformations")

    tab1, tab2, tab3 = st.tabs(["Grammar Agent", "Style Agent", "Clarity Agent"])

    with tab1:
        g_res = final_state["grammar_result"]
        st.markdown("**Output:**")
        st.write(g_res["corrected"])
        st.markdown(f"**Agent confidence:** `{g_res.get('confidence', 'N/A')}`")
        st.markdown("**Changes detected:**")
        if g_res.get("changes"):
            for c in g_res["changes"]:
                st.write(f"- `{c.get('from')}` → `{c.get('to')}` ({c.get('type')})")
        else:
            st.write("No explicit changes detected or diff too small.")

    with tab2:
        s_res = final_state["style_result"]
        st.markdown("**Output:**")
        st.write(s_res["styled"])
        st.markdown(f"**Agent confidence:** `{s_res.get('confidence', 'N/A')}`")
        st.markdown("**Changes detected:**")
        if s_res.get("changes"):
            for c in s_res["changes"]:
                st.write(f"- `{c.get('from')}` → `{c.get('to')}` ({c.get('type')})")
        else:
            st.write("No explicit changes detected or diff too small.")

    with tab3:
        c_res = final_state["clarity_result"]
        st.markdown("**Output:**")
        st.write(c_res["clarified"])
        st.markdown(f"**Agent confidence:** `{c_res.get('confidence', 'N/A')}`")
        st.markdown("**Changes detected:**")
        if c_res.get("changes"):
            for c in c_res["changes"]:
                st.write(f"- `{c.get('from')}` → `{c.get('to')}` ({c.get('type')})")
        else:
            st.write("No explicit changes detected or diff too small.")

    # --- Candidate scores table ---
    st.subheader("📊 Candidate Comparison (Reviewer Agent)")

    cand_rows = []
    for c in review["candidates"]:
        cand_rows.append({
            "Candidate": c["name"],
            "Score": c["score"],
            "Text": c["text"],
            "Notes": c["notes"]
        })

    st.dataframe(cand_rows)

elif run_button and not user_text.strip():
    st.warning("Please enter some text before running TextDoctor.")