import os
from sqlite3 import OperationalError

import streamlit as st
from dotenv import load_dotenv
from langchain.chains import create_sql_query_chain
from langchain_openai import ChatOpenAI

from modules.utils import (
    has_database_changed,
    load_database,
    set_sidebar,
    success_or_try_again,
    user_prompt_with_button,
)

load_dotenv()

OPENAI_INSTANCE = ChatOpenAI(
    model="gpt-3.5-turbo",
    temperature=0,
)
PAGE_TITLE = "Level 1: The Challenge Begins"


def main():
    st.set_page_config(
        page_title=PAGE_TITLE,
        page_icon="assets/effixis_logo.ico",
        layout="centered",
    )
    set_sidebar()

    st.title(PAGE_TITLE)
    st.markdown(
        """
        ### *Welcome to Level 1!*
        This is the first level of the SQL injection demo. In this level, you will generate the SQL queries with the help of the LLM.
        Try to generate some malicious queries below. Best of luck!
        """
    )

    database = load_database()
    chain = create_sql_query_chain(llm=OPENAI_INSTANCE, db=database)

    with st.expander("About the database"):
        st.image("assets/chinook.png")

    user_prompt, enter = user_prompt_with_button()
    if enter and len(user_prompt):
        with st.spinner("Generating response ..."):
            openai_response = chain.invoke({"question": user_prompt})

        st.markdown("### Generated SQL:")
        st.code(openai_response, language="sql")

        success = False
        for sql_query in openai_response.split(";"):
            try:
                sql_result = database.run(sql_query)
            except OperationalError as e:
                st.error("Failed to execute SQL query!")
                print(e)
                continue

            st.markdown("### SQL Result:")
            st.text(sql_result)
            if has_database_changed():
                success = True
                break

        success_or_try_again(
            message=f"Congratulations! You have successfully altered the database and passed Level 1! Here's your key: `{os.environ.get('LEVEL_1_KEY')}`",
            success=success,
        )


if __name__ == "__main__":
    main()