import streamlit as st
import pandas as pd
from groq_llms import LLMHandler
#from openrouter_llms import LLMHandler
import tempfile
import os
from dotenv import load_dotenv

load_dotenv()

# Initialize LLMHandler
llm_handler = LLMHandler()


def process_csv(file, user_prompt):
    """Read CSV, generate responses using LLMHandler, and return processed DataFrame."""
    df = pd.read_csv(file)
    responses = []

    for _, row in df.iterrows():
        try:
            response = llm_handler.generate_response(user_prompt, row.to_dict())
            responses.append(response)
        except Exception as e:
            responses.append(f"Error: {e}")

    df["Generated Text"] = responses
    return df


def initialize_session_state():
    """Initialize session state variables"""
    if 'prompt_creation_method' not in st.session_state:
        st.session_state.prompt_creation_method = None
    if 'current_step' not in st.session_state:
        st.session_state.current_step = 'choose_method'
    if 'context' not in st.session_state:
        st.session_state.context = ""
    if 'questions' not in st.session_state:
        st.session_state.questions = []
    if 'answers' not in st.session_state:
        st.session_state.answers = {}
    if 'multiselect_answers' not in st.session_state:
        st.session_state.multiselect_answers = {}
    if 'custom_options' not in st.session_state:
        st.session_state.custom_options = {}
    if 'final_prompt' not in st.session_state:
        st.session_state.final_prompt = ""
    if 'direct_prompt' not in st.session_state:
        st.session_state.direct_prompt = ""


def display_progress_tracker():
    """Display current progress and previous responses"""
    with st.expander("📋 View Progress", expanded=True):
        if st.session_state.prompt_creation_method:
            st.write(f"**Method chosen:** {st.session_state.prompt_creation_method.title()}")

        if st.session_state.context:
            st.write("**Initial Context:**")
            st.info(st.session_state.context)
            if st.button("Edit Context", key="edit_context"):
                st.session_state.current_step = 'initial_context'
                st.rerun()

        if st.session_state.answers:
            st.write("**Your Responses:**")
            for i, question in enumerate(st.session_state.questions):
                if i in st.session_state.multiselect_answers:
                    answers = ", ".join(st.session_state.multiselect_answers[i])
                    st.success(f"Q: {question['question']}\nA: {answers}")
                elif i in st.session_state.answers:
                    st.success(f"Q: {question['question']}\nA: {st.session_state.answers[i]}")
            if st.button("Edit Responses", key="edit_responses"):
                st.session_state.current_step = 'answer_questions'
                st.rerun()

        if st.session_state.direct_prompt:
            st.write("**Your Direct Prompt:**")
            st.info(st.session_state.direct_prompt)
            if st.button("Edit Prompt", key="edit_direct_prompt"):
                st.session_state.current_step = 'direct_prompt'
                st.rerun()

        if st.session_state.final_prompt:
            st.write("**Final Generated Prompt:**")
            st.info(st.session_state.final_prompt)
            if st.button("Edit Final Prompt", key="edit_final_prompt"):
                st.session_state.current_step = 'edit_prompt'
                st.rerun()


# Streamlit UI
st.set_page_config(page_title="Invite AI", page_icon="💬", layout="wide")

# Header
st.title("Invite AI")
st.markdown(
    """
    Welcome to the Invitation Generator! This tool helps you create personalized invitations using the power of AI.
    """
)

# Initialize session state
initialize_session_state()

# Display progress tracker (always visible)
display_progress_tracker()

# Sidebar with instructions
st.sidebar.title("Instructions")
st.sidebar.markdown(
    """
    ### Template Download
    [Click here to download the suggested CSV template](http://surl.li/ptvzzv) 📥
    ### Suggested Requirements
    - **Unique Identifier for each receiver**
    - **Name of the receiver**
    - **Designation/Job title of the receiver**
    - **Company/Organisation where the receiver works**
    - **Areas the receiver is interested in / has expertise in**
    - **Categorize receivers into groups**
    
    [Note: The above template is for your reference, you are free to submit your own data.] 
    """
)

# Main content area with steps
st.markdown("---")  # Separator between progress tracker and current step

if st.session_state.current_step == 'choose_method':
    st.subheader("Choose Your Prompt Creation Method")

    col1, col2 = st.columns(2)

    with col1:
        st.markdown("""
        ### Guided Prompt Builder
        - Step-by-step assistance
        - AI-generated questions
        - Structured approach
        """)
        if st.button("Use Guided Builder"):
            st.session_state.prompt_creation_method = 'guided'
            st.session_state.current_step = 'initial_context'
            st.rerun()

    with col2:
        st.markdown("""
        ### Direct Prompt Entry
        - Write your own prompt
        - Complete control
        - Quick setup
        """)
        if st.button("Use Direct Entry"):
            st.session_state.prompt_creation_method = 'direct'
            st.session_state.current_step = 'direct_prompt'
            st.rerun()

elif st.session_state.current_step == 'direct_prompt':
    st.subheader("Enter Your Prompt")
    st.markdown(
        "Write your complete prompt for generating invitations. Include all necessary details and requirements.")

    direct_prompt = st.text_area(
        "Your Prompt:",
        value=st.session_state.direct_prompt,
        placeholder="Example: Generate a professional invitation for a product launch...",
        height=200
    )

    col1, col2 = st.columns([1, 5])
    with col1:
        if st.button("← Back"):
            st.session_state.current_step = 'choose_method'
            st.rerun()
    with col2:
        if st.button("Continue →"):
            if direct_prompt:
                st.session_state.direct_prompt = direct_prompt
                st.session_state.final_prompt = direct_prompt
                st.session_state.current_step = 'upload_process'
                st.rerun()
            else:
                st.error("Please enter a prompt before continuing.")

elif st.session_state.prompt_creation_method == 'guided':
    if st.session_state.current_step == 'initial_context':
        st.subheader("Step 1: Provide Initial Context")
        st.markdown("Briefly describe what your invitation is about (e.g., 'Launching a new GPU product')")

        context = st.text_area(
            "Context:",
            value=st.session_state.context,
            placeholder="Example: Launching a new GPU product for AI and HPC applications",
            height=100
        )

        col1, col2 = st.columns([1, 5])
        with col1:
            if st.button("← Back"):
                st.session_state.current_step = 'choose_method'
                st.rerun()
        with col2:
            if st.button("Generate Questions →"):
                if context:
                    st.session_state.context = context
                    st.session_state.questions = llm_handler.generate_questions(context)
                    st.session_state.current_step = 'answer_questions'
                    st.rerun()
                else:
                    st.error("Please provide context before proceeding.")


    # In the answer_questions section of your code, replace the multiselect implementation with this:

    elif st.session_state.current_step == 'answer_questions':

        st.subheader("Step 2: Answer Questions")

        for i, question in enumerate(st.session_state.questions):

            if 'choices' in question:

                # Get previously selected options

                previous_selections = st.session_state.multiselect_answers.get(i, [])

                # Initialize base choices

                base_choices = question['choices'].copy()

                if "Custom" not in base_choices:
                    base_choices.append("Custom")

                # Add any previous custom value to the choices if it exists

                custom_values = [x for x in previous_selections if x not in question['choices'] and x != "Custom"]

                all_choices = base_choices + custom_values

                # Handle word count questions differently

                if any(word in question['question'].lower() for word in ['word count', 'words', 'length']):

                    selected_options = st.multiselect(

                        question['question'],

                        options=all_choices,

                        default=previous_selections,

                        key=f"multiselect_{i}"

                    )

                    if "Custom" in selected_options:

                        # Pre-fill with previous custom value if exists

                        default_custom = next((x for x in previous_selections if x not in base_choices), "")

                        custom_value = st.text_input(

                            "Enter custom word count:",

                            value=default_custom,

                            key=f"custom_{i}"

                        )

                        if custom_value:

                            try:

                                word_count = int(custom_value)

                                if word_count > 0:

                                    selected_options = [opt for opt in selected_options if opt != "Custom"]

                                    if str(word_count) not in selected_options:
                                        selected_options.append(str(word_count))

                                else:

                                    st.error("Please enter a positive number")

                            except ValueError:

                                st.error("Please enter a valid number")


                else:

                    # Regular non-numeric multiselect handling

                    selected_options = st.multiselect(

                        question['question'],

                        options=all_choices,

                        default=previous_selections,

                        key=f"multiselect_{i}"

                    )

                    if "Custom" in selected_options:

                        # Pre-fill with previous custom value if exists

                        default_custom = next((x for x in previous_selections if x not in base_choices), "")

                        custom_value = st.text_input(

                            "Enter your custom response:",

                            value=default_custom,

                            key=f"custom_{i}"

                        )

                        if custom_value:

                            selected_options = [opt for opt in selected_options if opt != "Custom"]

                            if custom_value not in selected_options:
                                selected_options.append(custom_value)

                # Update session state

                st.session_state.multiselect_answers[i] = selected_options

                st.session_state.answers[i] = ", ".join(selected_options) if selected_options else ""


            else:

                # Handle non-choice questions

                st.session_state.answers[i] = st.text_input(

                    question['question'],

                    value=st.session_state.answers.get(i, ""),

                    key=f"question_{i}"

                )

        col1, col2 = st.columns([1, 5])

        with col1:

            if st.button("← Back"):
                st.session_state.current_step = 'initial_context'

                st.rerun()

        with col2:

            if st.button("Generate Prompt →"):

                if all(st.session_state.answers.values()):

                    st.session_state.final_prompt = llm_handler.generate_final_prompt(

                        st.session_state.context,

                        st.session_state.questions,

                        st.session_state.answers

                    )

                    st.session_state.current_step = 'edit_prompt'

                    st.rerun()

                else:

                    st.error("Please answer all questions before proceeding.")
    elif st.session_state.current_step == 'edit_prompt':
        st.subheader("Step 3: Review and Edit Final Prompt")
        edited_prompt = st.text_area(
            "Edit your prompt if needed:",
            value=st.session_state.final_prompt,
            height=200
        )

        col1, col2 = st.columns([1, 5])
        with col1:
            if st.button("← Back"):
                st.session_state.current_step = 'answer_questions'
                st.rerun()
        with col2:
            if st.button("Continue to Upload →"):
                st.session_state.final_prompt = edited_prompt
                st.session_state.current_step = 'upload_process'
                st.rerun()

# Common upload and processing section for both paths
if st.session_state.current_step == 'upload_process':
    st.subheader("Upload and Process")
    uploaded_file = st.file_uploader("📂 Upload CSV File", type=["csv"])

    col1, col2 = st.columns([1, 5])
    with col1:
        if st.button("← Back"):
            if st.session_state.prompt_creation_method == 'guided':
                st.session_state.current_step = 'edit_prompt'
            else:
                st.session_state.current_step = 'direct_prompt'
            st.rerun()

    if uploaded_file is not None and st.session_state.final_prompt:
        st.write("⏳ Processing your file... Please wait.")
        processed_df = process_csv(uploaded_file, st.session_state.final_prompt)

        st.write("### Generated Invitations")
        st.dataframe(processed_df, use_container_width=True)

        with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file:
            processed_df.to_csv(temp_file.name, index=False)
            temp_file.close()

            st.download_button(
                label="📥 Download Results CSV",
                data=open(temp_file.name, "rb"),
                file_name="generated_invitations.csv",
                mime="text/csv",
            )
            os.unlink(temp_file.name)

# Reset button (moved to sidebar)
st.sidebar.markdown("---")
if st.sidebar.button("🔄 Start Over"):
    st.session_state.prompt_creation_method = None
    st.session_state.current_step = 'choose_method'
    st.session_state.context = ""
    st.session_state.questions = []
    st.session_state.answers = {}
    st.session_state.multiselect_answers = {}
    st.session_state.custom_options = {}
    st.session_state.final_prompt = ""
    st.session_state.direct_prompt = ""
    st.rerun()

st.markdown("---")
st.markdown("💡 **Tip:** Ensure your data aligns with the provided template for accurate results.")