import json
import security
import streamlit as st
from streamlit_extras.add_vertical_space import add_vertical_space
import os
import streamlit as st
import pandas as pd
from io import StringIO
from datetime import datetime
import boto3
import util
import requests
from sf_conn import get_secret, save_to_sf
from io import StringIO, BytesIO
import time
from sf_conn import get_client_names, insert_upload_logs
from constants import USER_LIST
from util import logger

(REDIRECT_URI, create_batch_url, doczy_pipeline) = util.load_page_details(5)

user_list = USER_LIST
if "uploading" not in st.session_state:
    st.session_state.uploading = False
if "upload_key" not in st.session_state:
    st.session_state.upload_key = 0
if "file_list" not in st.session_state:
    st.session_state.file_list = []
if "show_batchID" not in st.session_state:
    st.session_state.show_batchID = False
if "landing_zone" not in st.session_state:
    st.session_state.landing_zone = "contracts-landing-zone"
if "batch_id" not in st.session_state:
    st.session_state.batch_id = "failed_cases"
if "client_bucket" not in st.session_state:
    st.session_state.client_bucket = "default_bucket"


def insert_upload_logs(batch_id, client_name, file_name, upload_datetime, upload_user):
    """
    Input: batch_id, client_name, file_name, upload_datetime, upload_user
    Output: status of the insert query
    """
    try:
        return "Log inserted successfully"
    except Exception as e:
        return e


st.set_page_config(layout="wide")
# # # Sidebar contents
with st.sidebar:
    st.title("Doczy.AI ™")
    st.markdown(
        """
    ## About
    This app extracts data from contracts
 
    """
    )
    add_vertical_space(15)
    # st.write("Doczy")

# AARETE LOGO
x, y, z = st.columns([15, 2, 15])
with y:
    st.image("aaretelogo.png")

hide_img_fs = """
<style>
button[title="View fullscreen"]{
    visibility: hidden;}
</style>
"""
st.markdown(hide_img_fs, unsafe_allow_html=True)

_, c1 = st.columns([5, 1])
try:
    util.setup_page(REDIRECT_URI)
except Exception as e:
    st.write(f"SSO Failed = {e}")
    st.session_state["user_info"] = {
        "mail": "maamseek@aarete.com",
        "displayName": "Mayank Aamseek",
    }
try:
    c1.write(f"User: **{st.session_state.user_info['displayName']}**")
    user_mail = st.session_state.user_info["mail"]
except KeyError as e:
    st.write("Session Expired.")
    auth_url = security.get_auth_url(REDIRECT_URI)
    st.markdown(
        f"<a href='{auth_url}' target='_self'>Sign In</a>", unsafe_allow_html=True
    )
    st.stop()

s3_client = boto3.client(
    "s3",
    region_name="us-east-2",
)

client_list = [
    "doczy-ai-client-1",
    "Delaware First Health, Inc.",
    "Community Health Choice, Inc",
    "CareSource Network Partners LLC",
    "HealthNet of Cali",
    "Oklahoma Complete Health, Inc",
    "HealthFirst",
    "Molina Healthcare of TX",
    "AvMed",
    "Arizona Care1st",
    "WellCare New Jersey",
]

# This is the list of client fetched from Snowflake
# TODO: Need to update the streamlit code to use the client names from this list
# And use the s3 paths to save the objects for the respective client
client_list, s3_paths = get_client_names()
logger.info(f"client_list={client_list}")
logger.info(f"s3_paths={s3_paths}")
client_s3_paths = dict(zip(client_list, s3_paths))

client_row = st.columns([0.1, 0.8])
with client_row[0]:
    st.write("**Client Name**")
with client_row[1]:
    client = st.selectbox(
        "Client Name", (client_list), label_visibility="collapsed", index=None
    )

client_bucket = client_s3_paths.get(client)

# to be deleted when client buckets are created
# client_bucket = 'doczyai-use2-d-cn1-s3-textract-processing-001'

file_row = st.columns([0.1, 0.8])
with file_row[0]:
    st.write("**Upload Files**")
with file_row[1]:
    file_list = st.file_uploader(
        "Upload",
        type=["docx", "tiff", "pdf"],
        accept_multiple_files=True,
        label_visibility="collapsed",
        help="Only PDF, TIFF and DOCX file formats are supported.",
        disabled=st.session_state.uploading,
        key=st.session_state.upload_key,
    )

add_vertical_space(2)
df = pd.DataFrame(columns=["Contract Name"])
df["Contract Name"] = file_list
file_names = []
buttons = st.columns([0.4, 0.4, 0.2])


def set_uploading_state():
    if not client == None and not len(file_list) == 0:
        st.session_state.file_list = file_list
        st.session_state.upload_key += 1
        st.session_state.uploading = True


with buttons[1]:
    if st.button("Create Batch", on_click=set_uploading_state):
        file_list = st.session_state.file_list
        if client == None:
            st.error("No Client Name Selected.")
        elif len(file_list) == 0:
            st.error("No Files Selected.")
        else:
            with st.spinner("Running..."):
                myobj = {"client-bucket-name": client_bucket}
                response = requests.post(create_batch_url, json=myobj)
                if response.status_code >= 200 and response.status_code < 300:
                    try:
                        batch_id = json.loads(json.loads(response.text)["body"])[
                            "batch_id"
                        ]
                        landing_zone = json.loads(json.loads(response.text)["body"])[
                            "landing_zone"
                        ]
                        landing_zone = "contracts-landing-zone/"
                    except:
                        # st.write(myobj)
                        # st.write(response.text)
                        st.write("Internal Error. Reach out to Doczy.AI Team.")
                        batch_id = "failed_cases"
                        landing_zone = "contracts-landing-zone/"
                else:
                    st.error("Failed")
                for uploaded_file in file_list:
                    stringio = BytesIO(uploaded_file.getvalue())
                    stringio.seek(0)
                    s3_client.put_object(
                        Bucket=client_bucket,
                        Body=stringio.getvalue(),
                        Key=landing_zone + batch_id + "/" + str(uploaded_file.name),
                    )
                    s3_client.put_object_tagging(
                        Bucket=client_bucket,
                        Key=landing_zone + batch_id + "/" + str(uploaded_file.name),
                        Tagging={"TagSet": [{"Key": "BatchId", "Value": batch_id}]},
                    )

                    # TODO: Test this insert function with snowflake
                    upload_log = insert_upload_logs(
                        batch_id,
                        client,
                        str(uploaded_file.name),
                        datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                        user_mail,
                    )
                    st.write(upload_log)

                    file_names.append(str(uploaded_file.name))
                st.session_state.uploading = False
                st.session_state.show_batchID = True
                st.session_state.client_bucket = client_bucket
                st.session_state.batch_id = batch_id
                st.session_state.landing_zone = landing_zone
                st.rerun()
    if st.session_state.show_batchID:
        batch_id = st.session_state.batch_id
        client_bucket = st.session_state.client_bucket
        landing_zone = st.session_state.landing_zone
        st.write("Batch ID Created - ")
        st.code(f"{batch_id}")
        st.write(f"Files uploaded to s3://{client_bucket}/{landing_zone}{batch_id}")
        st.session_state.show_batchID = False
        st.session_state.file_list = []
        st.session_state.batch_id = "failed_cases"
        st.session_state.client_bucket = "default_bucket"
        st.session_state.landing_zone = "contracts-landing-zone"

    # @st.cache_data
    # def convert_df(df):
    #    return df.to_csv(index=False).encode('utf-8')

    # csv = convert_df(df)

    # df = df.reset_index()  # make sure indexes pair with number of rows

    # contract_list = []
    # for index, row in df.iterrows():
    #     group_list = []
    #     if row['Unique Key']:
    #         group_list.append('Unique Key')
    #     if row['Pricing Before Carveouts']:
    #         group_list.append('Pricing Before Carveouts')
    #     if row['Contract Related']:
    #         group_list.append('Contract Related')
    #     if row['Provider']:
    #         group_list.append('Provider')
    #     if row['Timeline']:
    #         group_list.append('Timeline')
    #     if row['Carveout Indicator']:
    #         group_list.append('Carveout Indicator')
    #     if row['Carveout Methodology']:
    #         group_list.append('Carveout Methodology')
    #     entry_dict = {
    #         "contract_name": row['Contract Name'],
    #         "groups": group_list,
    #         "contract_source_path": "batches/batch_1/"+client+"/"+row['Contract Name']
    #     }
    #     contract_list.append(entry_dict)
    # contract_list = list(df['Contract Name'])

    # myobj = {
    #   "s3_bucket": 'doczy-dev-infra-textract',
    #   "batch_id": "1",
    #   "client_name": client,
    #   "username": user_mail,
    #   "contract_list": contract_list
    # }

    # buttons = st.columns([0.8, 0.2])
    # with buttons[0]:
    #     st.download_button("Download Table", csv, "file.csv", "text/csv", key='download-csv')
    # with buttons[1]:
    #     if st.button("Upload to DB"):
    #         response = requests.post(doczy_pipeline, json = myobj)
    #         if response.status_code >= 200 and response.status_code < 300:
    #             st.write("Success")
    #         else:
    #             st.write("Failed")
    #         # st.write(response.text)