from huggingface_hub import Repository
from huggingface_hub import login
import os
from dotenv import load_dotenv

load_dotenv()

login(token = os.environ['HF_TOKEN'])

repo = Repository(
    local_dir="agent_function",
    repo_type="dataset",
    clone_from=os.environ['DATASET'],
    token=True
)
repo.git_pull()

import streamlit as st
import time
from agent_function.agent import answer,check_password

st.title("Tynox Chatbot")

if "messages" not in st.session_state:
    st.session_state.messages = []

if "auth" not in st.session_state:
    st.session_state.auth = False

if "chart" not in st.session_state:
    st.session_state.chart = None

if "report" not in st.session_state:
    st.session_state.report = None

if st.session_state.auth:
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if prompt := st.chat_input("What is up?"):
        st.chat_message("user").markdown(prompt)
        st.session_state.messages.append({"role": "user", "content": prompt})
        result=answer(prompt)
        assistant_message = st.chat_message("assistant")
        placeholder = assistant_message.empty()

        typed_response = ""
        for chunk in result:
            typed_response += chunk
            placeholder.markdown(typed_response)
            time.sleep(0.01) 
        
        if st.session_state.chart:
            st.chat_message("assistant").plotly_chart(
                st.session_state.chart
            )
            st.session_state.chart = None
        
        if st.session_state.report:
            def onclick():
                st.session_state.report = None
            st.chat_message("assistant").download_button(
                label=f"{st.session_state.report}",
                data=st.session_state.report_df,
                file_name=f"{st.session_state.report}_report.csv",
                mime="text/csv",
                on_click=onclick()
            )


        st.session_state.messages.append({"role": "assistant", "content": typed_response})
else:
    password = st.text_input("Password")
    if password:
        valid = check_password(password)
        if valid:
            st.session_state.auth = valid
            st.rerun()
        else:
            st.error('Incorrect password',icon='❌')