import streamlit as st
from src.main import ConversationalResponse
import os

# Constants
ROLE_USER = "user"
ROLE_ASSISTANT = "assistant"
# MAX_MESSAGES = 5  # Commented out but can be reactivated later

st.set_page_config(page_title="Chat with Git CodeLlama", page_icon="🦙")
st.title("Chat with Git CodeLlama 🦙📚")
st.markdown("by [Rohan Kataria](https://www.linkedin.com/in/imrohan/) view more at [VEW.AI](https://vew.ai/)")
st.markdown("This app allows you to chat with Git code files. You can paste link to the Git repository and ask questions about it. In the background uses the Git Loader and ConversationalRetrieval chain from langchain, Ollama, Streamlit for UI.")

@st.cache_resource(ttl="1h")
def load_agent(url, branch, file_filter):
    with st.spinner('Loading Git documents...'):
        agent = ConversationalResponse(url, branch, file_filter)
    st.success("Git Loaded Successfully")
    return agent

def main():

    git_link = st.sidebar.text_input("Enter your Git Link")
    branch = st.sidebar.text_input("Enter your Git Branch")
    file_filter = st.sidebar.text_input("Enter the Extension of Files to Load eg. py,sql,r (no spaces)")

    if "agent" not in st.session_state:
        st.session_state["agent"] = None
        st.session_state["user_message_count"] = 0

    if st.sidebar.button("Load Agent"):
        if git_link and branch and file_filter:
            try:
                st.session_state["agent"] = load_agent(git_link, branch, file_filter)
                st.session_state["messages"] = [{"role": ROLE_ASSISTANT, "content": "How can I help you?"}]
                st.session_state["user_message_count"] = 0
            except Exception as e:
                st.sidebar.error(f"Error loading Git repository: {str(e)}")
                return

    if st.session_state["agent"]:  # Chat will only appear if the agent is loaded
        for msg in st.session_state.messages:
            st.chat_message(msg["role"]).write(msg["content"])

        # if st.session_state["user_message_count"] < MAX_MESSAGES:  # Commented out to remove message limit check
        user_query = st.chat_input(placeholder="Ask me anything!")

        if user_query:
            st.session_state.messages.append({"role": ROLE_USER, "content": user_query})
            st.chat_message(ROLE_USER).write(user_query)
            st.session_state["user_message_count"] += 1

            # Generate the response
            with st.spinner("Generating response"):
                response = st.session_state["agent"](user_query)

            # Display the response immediately
            st.chat_message(ROLE_ASSISTANT).write(response)

            # Add the response to the message history
            st.session_state.messages.append({"role": ROLE_ASSISTANT, "content": response})
        # else:
            # st.warning("Your message limit is over. Contact [Rohan Kataria](https://www.linkedin.com/in/imrohan/) to increase the limit.")
            # Commented out the warning message about message limit

if __name__ == "__main__":
    main()