import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from io import BytesIO
import base64
import random
import io
import re
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from streamlit_tags import st_tags
from streamlit_vertical_slider import vertical_slider
import pdf_generator

# Set page config
st.set_page_config(
    page_title="Pilot Drafting",
    page_icon="🛫",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Password protection
def check_password():
    def password_entered():
        if st.session_state["password"] == st.secrets["app_password"]:
            st.session_state["password_correct"] = True
            del st.session_state["password"]
        else:
            st.session_state["password_correct"] = False

    if "password_correct" not in st.session_state:
        st.markdown("\n\n")
        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
        st.divider()
        st.info("Developed by Milan Mrdenovic © IBM Norway 2024")
        return False
    elif not st.session_state["password_correct"]:
        st.markdown("\n\n")
        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
        st.divider()
        st.info("Developed by Milan Mrdenovic © IBM Norway 2024")
        st.error("😕 Password incorrect")
        return False
    else:
        return True

if not check_password():
    st.stop()

# Initialize session state
if 'current_page' not in st.session_state:
    st.session_state.current_page = 0

if 'answers' not in st.session_state:
    st.session_state.answers = {
        'workload_scope': {
            'time_commitments': '',
            'deadlines': '',
            'feature_prioritization': ''
        },
        'useful_assets': {
            'assets': []
        },
        'pilot_evaluation': {
            'metrics': []
        }
    }

# Define the content for each page
pages = [
    {
        'title': "Scope a Manageable Workload",
        'content': """
        Define the scope of your project to ensure a manageable workload. Consider the time commitments, deadlines, and feature prioritization.
        """,
        'input_key': 'workload_scope',
        'input_type': 'custom'
    },
    {
        'title': "Useful Assets and Libraries",
        'content': """
        Discuss the useful assets and libraries that can be utilized for the project. Consider Langchain, Langflow, ChromaDB, Jupyter Notebooks, etc.
        """,
        'input_key': 'useful_assets',
        'input_type': 'custom'
    },
    {
        'title': "Pilot Evaluation",
        'content': """
        Plan for the metrics that need to be tracked for the project. Consider queries, token usage, F1 scores, precision, etc.
        """,
        'input_key': 'pilot_evaluation',
        'input_type': 'custom'
    },
    {
        'title': "Generate Evaluation Report",
        'content': "You have completed the Pilot Drafting. \nClick the button below to generate and download your PDF report.",
        'input_key': None
    }
]

st.session_state.pages = pages

# Main Streamlit app
st.title("Pilot Drafting")

# Navigation buttons
col1, col2, col3 = st.columns([1, 2, 1])
with col1:
    if st.session_state.current_page > 0:
        if st.button("Back"):
            st.session_state.current_page -= 1
            st.rerun()
with col3:
    if st.session_state.current_page < len(pages) - 1:
        if st.button("Next", use_container_width=True):
            st.session_state.current_page += 1
            st.rerun()

# Display current page
current_page = pages[st.session_state.current_page]
st.header(current_page['title'])

with st.expander("Description", expanded=False):
    st.markdown(current_page['content'])

# Input fields
if 'input_key' in current_page and current_page['input_key'] is not None:
    if current_page['input_key'] == 'workload_scope':
        st.subheader("Workload Scope")
        col1, col2, col3 = st.columns(3)
        with col1:
            st.session_state.answers['workload_scope']['time_commitments'] = st.text_area(
                "Time Commitments:",
                value=st.session_state.answers['workload_scope'].get('time_commitments', ""),
                key="time_commitments",
                height=150
            )
        with col2:
            st.session_state.answers['workload_scope']['deadlines'] = st.text_area(
                "Deadlines:",
                value=st.session_state.answers['workload_scope'].get('deadlines', ""),
                key="deadlines",
                height=150
            )
        with col3:
            st.session_state.answers['workload_scope']['feature_prioritization'] = st.text_area(
                "Feature Prioritization:",
                value=st.session_state.answers['workload_scope'].get('feature_prioritization', ""),
                key="feature_prioritization",
                height=150
            )

    elif current_page['input_key'] == 'useful_assets':
        st.subheader("Useful Assets and Libraries")
        assets = st_tags(
            label='Enter Useful Assets and Libraries:',
            text='Press enter to add more',
            value=st.session_state.answers['useful_assets'].get('assets', []),
            suggestions=[],
            maxtags=5,
            key='useful_assets'
        )
        st.session_state.answers['useful_assets']['assets'] = assets

    elif current_page['input_key'] == 'pilot_evaluation':
        st.subheader("Pilot Evaluation")
        metrics = st_tags(
            label='Enter Metrics to Track:',
            text='Press enter to add more',
            value=st.session_state.answers['pilot_evaluation'].get('metrics', []),
            suggestions=[],
            maxtags=5,
            key='pilot_evaluation'
        )
        st.session_state.answers['pilot_evaluation']['metrics'] = metrics

# Generate PDF button (only on the last page)
if st.session_state.current_page == len(pages) - 1:
    if st.button("Generate and Download PDF", use_container_width=True):
        pdf = pdf_generator.generate_pdf(st.session_state)
        st.download_button(
            label="Download PDF",
            data=pdf,
            file_name="Pilot_Drafting.pdf",
            mime="application/pdf",
            use_container_width=True
        )

# Display progress
st.progress((st.session_state.current_page + 1) / len(pages))

st.divider()