from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from langchain_aws import ChatBedrock
import streamlit as st
import os
from dotenv import load_dotenv

load_dotenv()

class ExcelAnalyser:
    def __init__(self, uri):  # Fix the method name
        self.uri = uri
    
    def connect_to_uri(self):
        db = SQLDatabase.from_uri(self.uri) 
        llm = ChatBedrock(
        model="anthropic.claude-3-5-sonnet-20240620-v1:0",
     model_kwargs={
        "temperature": 0,
    },
    region='us-east-1',
        aws_access_key_id=os.getenv('aws_access_key'),
        aws_secret_access_key=os.getenv('aws_secret_key')
        )
        agent_executor = create_sql_agent(llm, db=db, verbose=True)
        return agent_executor

def chat_interface():
    st.title("Chat with your Excel Data")
     # Add debug info
     

    # Check if database path exists in session state
    if 'db_path' not in st.session_state:
        st.warning("Please upload an Excel file first!")
        return
    db = SQLDatabase.from_uri(st.session_state['db_path'])
    tables = db.get_usable_table_names()
    st.write(f"Available tables: {tables}")
    # Initialize chat history
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    # Display chat history
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Accept user input
    if prompt := st.chat_input("Ask questions about your Excel data"):
        # Display user message
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        with st.chat_message("assistant"):
            try:
                analyser = ExcelAnalyser(st.session_state['db_path'])
                agent = analyser.connect_to_uri()
                response = agent.invoke(prompt)
                st.markdown(response['output'])
                st.session_state.messages.append({"role": "assistant", "content": response['output']})
            except Exception as e:
                error_message = f"Error: {str(e)}"
                st.error(error_message)
                st.session_state.messages.append({"role": "assistant", "content": error_message})

if __name__ == '__main__':
    chat_interface()