import streamlit as st
from api import fetch_okrs
from config import INPUT_TEMPLATE, team
from utils import extract_llm_response

def proposer_page():
    # Streamlit App Layout
    st.title("AO/PM OKR Proposer")

    # Input Section and Buttons Row
    st.subheader("Enter your idea or goal:")
    user_input = st.text_area(
        "Input your idea here:",
        value=st.session_state.get("user_input", INPUT_TEMPLATE.strip()),
        height=300,
    )

    generate_okrs_clicked = st.button("Generate OKR Proposal")

    if generate_okrs_clicked:
        if not user_input.strip():
            st.warning("Please provide some input before generating OKRs.")
        else:
            with st.spinner("Generating OKRs..."):
                response = fetch_okrs(user_input=user_input)
                if response:
                    # Extract Objective and Key Results from response
                    objective, key_results, hint = extract_llm_response(response)
                    st.session_state["objective"] = objective
                    st.session_state["key_results"] = key_results
                    st.session_state["hint"] = hint

    # Display Results Only if an OKR Has Been Generated
    if "objective" in st.session_state and "key_results" in st.session_state:
        # Ensure team members exist in session state
        if "team_members" not in st.session_state:
            # Sorting the list alphabetically by the first name
            sorted_names = sorted(team, key=lambda name: name.split()[0])
            st.session_state["team_members"] = sorted_names  # Default team members

        team_members = st.session_state["team_members"]

        # Display Objective Field with Responsibles Input Below It
        st.subheader("Proposal Objective:")
        objective_text = st.text_area(
            "Proposal Objective:",
            value=st.session_state.get("objective", ""),
            height=100,
        )

        responsible_for_objective = st.multiselect(
            "Select Responsibles for Objective:",
            options=team_members,
            default=[],
            key="responsibles_objective"
        )

        # Display Key Results with Responsibles Below Each One
        st.subheader("Proposal Key Results:")
        key_result_boxes = []
        responsibles_for_key_results = []

        for i, kr in enumerate(st.session_state["key_results"], start=1):
            kr_text = st.text_area(f"Key Result {i}:", value=kr, key=f"kr_{i}")
            responsible_for_kr = st.multiselect(
                f"Select Responsibles for Key Result {i}:",
                options=team_members,
                default=[],
                key=f"responsibles_kr_{i}"
            )
            
            key_result_boxes.append(kr_text)
            responsibles_for_key_results.append(responsible_for_kr)

        # Finalize Button Center-Aligned
        if st.button("Finalize"):
            finalized_objective = objective_text.strip()
            finalized_key_results = [st.session_state[f"kr_{i+1}"].strip() for i in range(len(key_result_boxes))]

            # Append initials of responsibles to Objective and Key Results
            responsibles_list_objective = responsible_for_objective
            initials_objective = [f"[{''.join([part[0] for part in name.split()]).upper()}]" for name in responsibles_list_objective]
            initials_str_objective = ", ".join(initials_objective)

            finalized_objective = f"{initials_str_objective} {finalized_objective}"

            finalized_key_results_with_initials = []
            for i, kr in enumerate(finalized_key_results):
                responsibles_list_kr = responsibles_for_key_results[i]
                initials_kr = [f"{''.join([part[0] for part in name.split()]).upper()}" for name in responsibles_list_kr]
                initials_str_kr = ", ".join(initials_kr)
                finalized_key_results_with_initials.append(f"KR{i+1}: [{initials_str_kr}] {kr}")

            # Display finalized data in non-editable format (full width)
            st.subheader("Finalized Objective:")
            st.code(
                body=finalized_objective,
                language=None,
                wrap_lines=True
            )

            st.subheader("Finalized Key Results:")
            for kr in finalized_key_results_with_initials:
                with st.container():
                    col1, col2 = st.columns(spec=[0.1, 0.9])
                    with col1:
                        st.caption(kr.split(':')[0])
                    
                    with col2:
                        st.code(
                            body=kr.split(':')[-1],
                            language=None,
                            wrap_lines=True
                        )