# app.py
import os
import json
import streamlit as st
from app_assets import app_logo
import io
from app_utils import parse_data_dictionary, display_notebook, style_css
from workflow_nodes import (
    think_sections_node,
    enhance_plan_node,
    write_code_node,
    modify_notebook_node,
    execute_notebook_node,
    correct_notebook_node,
    write_insights_node
)
from workflow_nodes import InteractiveCaseStudyState

st.set_page_config(layout="wide")

# Apply CSS styling
# This will apply the styles defined in style.css to the Streamlit app
style_css("/home/user/app/src/style.css")

st.title("📓 NBforge")

#st.logo(app_logo, size='large', link='https://www.mygreatlearning.com/')


# ─── Session State Initialization ─────────────────────────────────────────────
if "state" not in st.session_state:
    st.session_state.state = InteractiveCaseStudyState()
    st.session_state.plan_ready = False
    st.session_state.plan_approved = None
    st.session_state.notebook_ready = False
    st.session_state.notebook_approved = None
    st.session_state.exec_ready = False
    st.session_state.execution_approved = None

s = st.session_state.state

# ─── 1) Collect Inputs ────────────────────────────────────────────────────────
st.header("1. 🔧 Provide Your Inputs")
s["domain"] = st.selectbox("📂 Select Domain", [
    "EDA", "Machine Learning", "Deep Learning", "Natural Language Processing",
    "Computer Vision", "Generative AI Prompt Engineering", "Generative AI RAG"
])
s["topics_and_subtopics"] = st.text_area("📚 Topics & Subtopics", height=80)
s["problem_statement"] = st.text_area("🧩 Problem Statement", height=80)
s["dataset_type"] = st.selectbox("📁 Dataset Type", ["csv", "pdf", "images", "json"])

uploaded = st.file_uploader("⬆️ Upload Dataset", accept_multiple_files=False)
if uploaded:
    os.makedirs("data", exist_ok=True)
    path = f"data/{uploaded.name}"
    with open(path, "wb") as f:
        f.write(uploaded.getbuffer())
    s["dataset_file_path"] = path
    st.success(f"Dataset saved to `{path}`")

s["data_dictionary"] = {}
dd = st.text_area("📖 Data Dictionary")
s["data_dictionary"] = parse_data_dictionary(dd)

s["additional_instructions"] = st.text_area(
    "📝 Additional Instructions (optional)", height=68
)

# Kick off plan generation
if not st.session_state.plan_ready:
    if st.button("🚀 Generate Plan") and uploaded:
        with st.spinner("Designing notebook plan…"):
            s = think_sections_node(s)
        st.session_state.plan_ready = True
        st.session_state.plan_approved = None
        st.session_state.notebook_ready = False
        st.session_state.exec_ready = False

# ─── 2) Review & Enhance Plan ────────────────────────────────────────────────
if st.session_state.plan_ready:
    st.subheader("2. 🔍 Review Notebook Plan")
    
    with st.expander("📑 View Generated Plan"):
        st.json(s["plan"])
    
    # ask approval once
    if st.session_state.plan_approved is None:
        st.session_state.plan_approved = st.radio("Approve this plan?", index=None, options=["YES", "NO"])
        st.write(f"Plan approved: {st.session_state.plan_approved}")

    
    # on approval, let user move on
    if st.session_state.plan_approved == "YES":
        if st.button("➡️ Next: Write Code"):
            with st.spinner("Writing notebook skeleton…"):
                s = write_code_node(s)
            st.session_state.notebook_ready = True
            st.session_state.notebook_approved = None

    # on rejection, collect feedback & loop back
    elif st.session_state.plan_approved == "NO":
        st.text_area("What should be changed in the plan?", key="plan_fb")
        st.write(f"Plan feedback: {st.session_state.plan_fb}")
        if st.button("🛠 Apply Plan Changes"):
            s["plan_feedback"] = st.session_state.plan_fb
            with st.spinner("Modifying notebook plan…"):
                s = enhance_plan_node(s)
            
            # reset approval so we review again
            st.session_state.plan_approved = None
            st.rerun()

# ─── 3) Review & Modify Notebook ─────────────────────────────────────────────
if st.session_state.notebook_ready:
    st.subheader("3. 💻 Review Notebook Skeleton")
    with st.expander("📄 View Notebook Skeleton"):
        display_notebook(io.StringIO(json.dumps(s["raw_notebook"], indent=2)))
        
        st.info("----------------------------------------------------------------")
        st.write("**NOTE**: If it’s hard to read here, download and open in Colab:")    
        st.download_button(
            "Download raw_notebook.ipynb",
            data=json.dumps(s["raw_notebook"], indent=2),
            file_name="raw_notebook.ipynb",
            mime="application/json",
        )

    if st.session_state.notebook_approved is None:
        st.session_state.notebook_approved = st.radio("Looks good?", index=None, options=["YES", "NO"])
        st.write(f"notebook approved: {st.session_state.notebook_approved}")

    if st.session_state.notebook_approved == "YES":
        if st.button("➡️ Next: Execute Notebook"):
            with st.spinner("Executing notebook…"):
                s = execute_notebook_node(s)
            st.session_state.exec_ready = True
            st.session_state.execution_approved = None

    elif st.session_state.notebook_approved == "NO":
        st.write("🛠 Please provide your feedback on the notebook skeleton:")
        st.text_area("What needs changing in the notebook skeleton?", key="nb_fb")
        st.write(f"Notebook feedback: {st.session_state.nb_fb}")
        if st.button("🛠 Apply Notebook Changes"):
            s["notebook_feedback"] = st.session_state.nb_fb
            with st.spinner("Modifying notebook skeleton…"):
                s = modify_notebook_node(s)
            st.session_state.notebook_approved = None
            st.rerun()



# ─── 4) Review & Correct Execution ───────────────────────────────────────────
if st.session_state.exec_ready:
    st.subheader("4. 🚀 Execution Results")
    with st.expander("📄 View Executed Notebook"):
        display_notebook(io.StringIO(s["executed_notebook"]))

        st.info("----------------------------------------------------------------")
        st.write("**NOTE**: If it’s hard to read here, download and open in Colab:")    
        st.download_button(
            "Download executed_notebook.ipynb",
            data=s["executed_notebook"],
            file_name="executed_notebook.ipynb",
            mime="application/json",
        )

    if s.get("execution_error"):
        # Always reset approval until they choose Download As-Is
        st.session_state.execution_approved = False
        st.error(f"❌ Execution failed:\n{s['execution_error']}")

        # Three-way branching
        st.session_state.exec_action = st.radio(
            "What would you like to do with this broken notebook?", index=None,
            options=["Download As-Is", "Provide Feedback", "Auto-Correct & Retry"]
        )
        st.write(f"Execution action: {st.session_state.exec_action}")

        if st.session_state.exec_action == "Download As-Is":
            # mark approved so we fall through to downloads
            st.session_state.execution_approved = True

        elif st.session_state.exec_action == "Provide Feedback":
            fb = st.text_area("📝 Describe how I should fix the error:", height=80, key="exec_feedback")
            st.write(f"Execution feedback: {st.session_state.exec_feedback}")
            if st.button("🛠 Apply Feedback"):
                s["execution_feedback"] = fb
                with st.spinner("Applying your feedback and re-executing…"):
                    s = correct_notebook_node(s)
                    s = execute_notebook_node(s)
                # loop back into this block
                st.session_state.exec_ready = True
                st.session_state.execution_approved = None
                st.rerun()

        elif st.session_state.exec_action == "Auto-Correct & Retry":
            if st.button("🔧 Auto-Correct & Retry"):
                s["execution_feedback"] = ""
                with st.spinner("Auto-correcting and re-executing…"):
                    s = correct_notebook_node(s)
                    s = execute_notebook_node(s)
                # loop back into this block
                st.session_state.exec_ready = True
                st.session_state.execution_approved = None
                st.rerun()

    # If execution succeeded OR user chose “Download As-Is”, show downloads
    if not s.get("execution_error") or st.session_state.execution_approved:
        st.success("✅ Ready to download your artifacts!")
        st.download_button("📥 Download plan.json", 
                           data=json.dumps(s["plan"], indent=2), 
                           file_name="plan.json", 
                           mime="application/json"
                           )
        st.download_button("📥 Download raw_notebook.ipynb", 
                           data=json.dumps(s["raw_notebook"], indent=2), 
                           file_name="raw_notebook.ipynb", 
                           mime="application/json"
                           )
        st.download_button("📥 Download executed_notebook.ipynb", 
                           data=s["executed_notebook"], 
                           file_name="executed_notebook.ipynb", 
                           mime="application/json"
                           )



# ─── 5) Generate Insights (Optional) ──────────────────────────────────────────
# Only show this once the user has downloaded (or chosen) the executed notebook
if st.session_state.exec_ready and (
    not s.get("execution_error") or st.session_state.execution_approved
):
    st.markdown("---")
    st.subheader("5. ✍️ Generate Insights (Optional)")

    # If there was an error and they chose Download As-Is, still allow insights
    if s.get("execution_error"):
        st.info("Your notebook contains errors, but you can still generate insights on available outputs.")
    else:
        st.success("Notebook executed end-to-end without errors!")
        st.info("You can now generate insights and conclusions based on the executed notebook.")
    
    # Ask permission
    st.session_state.insights_opt = st.radio(
        "Shall I proceed to write the insights and conclusion for this notebook?", index=None,
        options=["YES", "NO"]
    )

    if st.session_state.insights_opt == "YES":
        if st.button("📝 Generate Insights Now"):
            with st.spinner("Writing observations & final recommendations…"):
                s = write_insights_node(s)
            st.success("🎉 Insights generated!")
            
            # Render the final notebook with insights inline
            st.markdown("#### Final Notebook with Insights")
            with st.expander("📄 View Final Notebook with Insights"):
                display_notebook(io.StringIO(json.dumps(s["final_notebook"], indent=2)))
            
            # Offer download     
            st.download_button(
                "📥 Download Final Notebook (.ipynb)",
                data=json.dumps(s["final_notebook"], indent=2),
                file_name="final_notebook.ipynb",
                mime="application/json",
            )
            st.balloons()
            st.write("Thanks for using NBforge! 🚀")
    elif st.session_state.insights_opt == "NO":
        st.snow()
        st.write("👍 No problem! Feel free to revisit this notebook later or download any artifact above.")