import streamlit as st
from api import fetch_okrs
#from config import INPUT_TEMPLATE, team, cycle_definitions
from config.config import INPUT_EXAMPLE, cycle_definitions, team
from utils import extract_llm_response

import json

from ui.parameters import section_parameters

def section_user_input():
    # section for input parameters
    section_parameters()

    # 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_EXAMPLE.strip()),
        height=300,
    )

    st.session_state['user_input'] = user_input

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

    if generate_okrs_clicked:
        print("session_state:", st.session_state)
        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,
                    okr_cycle=st.session_state['okr_cycle'],
                    num_key_results=st.session_state['num_key_results'])
                
                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

def section_okr_proposals():
    # 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",
            max_selections=1
        )

        # 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(
                label=f'KR{i}',
                value=kr,
                key=f"proposal_kr_{i}"
            )
            #kr_text = st.text_area(f"Key Result {i}:", value=kr, key=f"kr_{i}")
            responsible_for_kr = st.pills(
                label=f"Select Responsibles for Key Result {i}:",
                options=team_members,
                default=[],
                key=f"responsibles_kr_{i}",
                selection_mode='multi'
            )
            
            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"proposal_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}"
            st.session_state['finalized_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}")
            st.session_state['finalized_key_results_with_initials'] = finalized_key_results_with_initials

            # 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
                        )
            export_data_as_json()

def export_data_as_json():
    # Prepare data for export
    finalized_key_results_with_initials = st.session_state.get('finalized_key_results_with_initials')
    export_data = {
        "user_input": st.session_state.get('user_input'),
        "finalized_objective": st.session_state.get('finalized_objective'),
        "key_results": [kr.split(':')[-1].strip() for kr in finalized_key_results_with_initials]
    }

    # Convert data to JSON string
    json_data = json.dumps(export_data, indent=4)

    # Add a download button
    st.download_button(
        label="Download Finalized OKR",
        data=json_data,
        file_name="okr_data.json",
        mime="application/json"
    )

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

    section_user_input()

    section_okr_proposals()