import streamlit as st
import pandas as pd
from utils.sql_utils import create_schema

MAX_FILES = 5

if "uploaded_dataframes" not in st.session_state:
    st.session_state.uploaded_dataframes = {}

st.header("Upload your CSV or Excel files")
st.caption("Maximum 5 files can be uploaded.")

num_uploaded_files = len(st.session_state.uploaded_dataframes)

disabled = num_uploaded_files >= MAX_FILES

uploaded_files = st.file_uploader(
    "Choose files",
    type=["csv", "xlsx"],
    accept_multiple_files=True,
    disabled=disabled,
)

if uploaded_files and not disabled:
    uploaded_count = 0
    for uploaded_file in uploaded_files:
        if len(st.session_state.uploaded_dataframes) < MAX_FILES:
            df = None
            try:
                if uploaded_file.name.endswith(".csv"):
                    df = pd.read_csv(uploaded_file)
                elif uploaded_file.name.endswith(".xlsx"):
                    df = pd.read_excel(uploaded_file)

                if uploaded_file.name in st.session_state.uploaded_dataframes:
                    st.toast(
                        f"File {uploaded_file.name} already uploaded. Skipping...",
                        icon="⚠️",
                    )
                else:
                    key = (
                        uploaded_file.name.lower()
                        .strip()
                        .replace(" ", "_")
                        .replace(".csv", "")
                        .replace(".xlsx", "")
                    )
                    st.session_state.uploaded_dataframes[key] = df
                    uploaded_count += 1
                print(f"Uploaded file: {str(uploaded_file.name)}")
            except Exception as e:
                st.error(f"Error reading file {uploaded_file.name}: {e}")
        else:
            st.warning(
                f"Maximum number of files({MAX_FILES}) reached. Cannot upload more files."
            )
            break
    if uploaded_count > 0:
        st.success(
            f"{uploaded_count} File(s) uploaded successfully. Total {len(st.session_state.uploaded_dataframes)} File(s)"
        )


if len(st.session_state.uploaded_dataframes) >= MAX_FILES:
    st.warning(
        f"Maximum number of files({MAX_FILES}) reached. Cannot upload more files."
    )


dataframes = st.session_state.uploaded_dataframes

if dataframes:
    st.header("Schema of Uploaded Files📚")
    schema = create_schema(dataframes)
    if "uploaded_df_schema" not in st.session_state:
        st.session_state.uploaded_df_schema = schema
    pretty_schema, markdown = st.tabs(["Schema", "Copy Schema in Markdown"])
    with pretty_schema:
        st.info(
            "You can copy this schema, and give it to any state of the art LLM models like (Gemini /ChatGPT /Claude etc) to cross check your answers.\n You can run the queries directly here, by using ***Manual Query Executer*** in the sidebar and download your results 😊",
            icon="ℹ️",
        )
        st.markdown(schema, unsafe_allow_html=True)
    with markdown:
        st.info(
            "You can copy this schema, and give it to any state of the art LLM models like (Gemini /ChatGPT /Claude etc) to cross check your answers.\n You can run the queries directly here, by using ***Manual Query Executer*** in the sidebar and download your results 😊",
            icon="ℹ️",
        )
        st.markdown(f"```\n{schema}\n```")