import streamlit as st
from merged import load_model_and_tokenizer, generate_code_stream
from pdf_utils import test_case

# ---------------- App Boot ----------------
st.set_page_config(page_title="Code + Test Case Generator", layout="wide")
st.write("🚀 App started")

# ---------------- Load model ONCE ----------------
@st.cache_resource
def load_model():
    return load_model_and_tokenizer()

tokenizer, model = load_model()

# ---------------- Session State ----------------
for key in ["selected_language", "generated_code", "pdf_bytes", "user_input"]:
    if key not in st.session_state:
        st.session_state[key] = "" if key != "selected_language" else None

# ---------------- UI Helpers ----------------
def select_language(lang):
    st.session_state.selected_language = lang
    st.session_state.generated_code = ""
    st.session_state.pdf_bytes = b""
    st.session_state.user_input = ""

def reset():
    for k in st.session_state:
        st.session_state[k] = "" if k != "selected_language" else None

# ---------------- UI ----------------
st.title("🧠 Generate Code & Download Test Cases")

# -------- Language Selection --------
if st.session_state.selected_language is None:
    st.subheader("Select Language")
    cols = st.columns(4)
    cols[0].button("Java", on_click=select_language, args=("Java",))
    cols[1].button("React.js", on_click=select_language, args=("React.js",))
    cols[2].button("Python", on_click=select_language, args=("Python",))
    cols[3].button("C++", on_click=select_language, args=("C++",))

# -------- Main App --------
else:
    lang = st.session_state.selected_language
    st.subheader(f"Selected Language: {lang}")
    st.button("🔄 Reset", on_click=reset)

    st.session_state.user_input = st.text_input(
        "Describe the task",
        value=st.session_state.user_input
    )

    # -------- Generate Code Button --------
    if st.button("⚙️ Generate Code"):
        if st.session_state.user_input.strip():

            st.session_state.generated_code = ""
            st.session_state.pdf_bytes = b""

            st.subheader("Generated Code")
            code_placeholder = st.empty()

            # 1️⃣ Code Generation Loader
            with st.spinner("⚙️ Generating code..."):
                for token in generate_code_stream(
                    lang,
                    st.session_state.user_input,
                    tokenizer,
                    model
                ):
                    st.session_state.generated_code += token
                    code_placeholder.code(st.session_state.generated_code)

            # 2️⃣ PDF Generation Loader
            with st.spinner("📄 Generating test cases PDF..."):
                st.session_state.pdf_bytes = test_case(
                    st.session_state.generated_code
                )

            st.success("✅ Test cases PDF generated")

        else:
            st.warning("Please enter a task")

    st.divider()

    # -------- Download PDF --------
    if st.session_state.get("pdf_bytes"):
        st.download_button(
            label="📥 Download Test Cases PDF",
            data=st.session_state.pdf_bytes,
            file_name="test_cases.pdf",
            mime="application/pdf",
            key="download_pdf"
        )