import json

import boto3
from langchain.prompts import PromptTemplate
from langchain.embeddings.bedrock import BedrockEmbeddings
from langchain.llms.bedrock import Bedrock
from langchain_community.vectorstores import Chroma
from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME, SOURCE_DIRECTORY
from langchain.chains import RetrievalQA

import streamlit as st
from streamlit_extras.add_vertical_space import add_vertical_space
import os
import pandas as pd
import numpy as np
import util
import anthropic
from pydantic import BaseModel
from typing import List
import re

REDIRECT_URI = 'https://doczy.aarete.com:8502'
user_list = ['maamseek@aarete.com', 'smahdavian@aarete.com', 'ahinge@aarete.com', 'akadam@aarete.com', 'pkatariya@aarete.com'
, 'piragavarapu@aarete.com', 'umistry@aarete.com', 'ahutchison@aarete.com', 'bgrunst@aarete.com', 'ddimeglio@aarete.com'
, 'vnair@aarete.com', 'kminhas@aarete.com','dculotta@aarete.com','cbull@aarete.com','sclark@aarete.com', 'fmohiuddin@aarete.com']

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")

_,c1= st.columns([5,1])
try:
    util.setup_page(REDIRECT_URI)
except:
    st.write("SSO Failed")
    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.")
    st.stop()

try:
    sf_secrets = json.loads(get_secret())
    conn = snowflake.connector.connect(
        user=sf_secrets.get('user'),
        password=sf_secrets.get('password'),
        account="aarete-doczyai",
        role = "DEVADMIN",
        warehouse="DEV_XS",
        database="DOCZY_DEV",
        schema="STG"
        )
    cur = conn.cursor()
    query = 'select * from "TRAINING_DATA_RAW"'
    cur.execute(query)
    field_values = pd.DataFrame.from_records(iter(cur), columns=[x[0] for x in cur.description])
    # st.write(field_values)
    field_values['Document_Name'] = field_values['DOCUMENT_NAME']
    # field_values['Contract ID'] = field_values['CONTRACT_TITLE']
    # error('table values are incorrect')
except:
    field_values = pd.read_csv('contract_field_values.csv', encoding='unicode_escape', skipinitialspace=True)
    field_values.rename(columns={'(internal) Document Name': 'Document_Name'}, inplace = True)
    # field_values.rename(columns={'(Internal) Carveout ID': 'Contract ID'}, inplace = True)
    # st.write("conn failed")

try:
    query = 'select * from "PROMPT_CONFIG"'
    cur.execute(query)
    fields = pd.DataFrame(cur.fetchall())
    fields.rename(columns={'FIELD_DESC': 'Field Name'}, inplace = True)
    fields.rename(columns={'PROMPT': 'Interrogation Question?'}, inplace = True)
    fields.rename(columns={'GROUP_ID': 'PRIORITY'}, inplace = True)
    fields.rename(columns={'FIELD_NAME': 'SF_DB_COL_NAME'}, inplace = True)
    fields.rename(columns={'FM_MODEL_ID': 'llm_selected'}, inplace = True)
    error('table is empty')
except:
    fields = pd.read_csv('contract_fields.csv', encoding='unicode_escape', skipinitialspace=True)
    fields = fields.drop_duplicates(subset='Field Name', keep="first").sort_values('Field Name')
    fields = fields[~fields['Field Name'].isnull()]

s3_client = boto3.client('s3',
    region_name="us-east-2"
)
bucket = 'doczy-dev-infra-textract'
objects = s3_client.list_objects_v2(Bucket=bucket, Prefix="training-data/contract-text-file/")
file_list = []
for obj in objects['Contents']:
    if not obj['Key'].endswith('/'):
        file_list.append(obj['Key'])

contract_list = sorted(file_list)

# to be deleted later
contract_list = [contract for contract in contract_list if contract.rsplit('/',1)[1].replace(' MU','').replace('_MU','').replace('.txt','') in list(field_values['Document_Name'])]

if user_mail in user_list: 

    file_row = st.columns([0.2, 0.7, 0.1])
    with file_row[0]:
        st.write("**Contract Name**")
    with file_row[1]:
        # file_name = st.text_input("**Contract Name**", label_visibility = "collapsed")
        file_name = st.selectbox('Select a file', contract_list + ['All'], label_visibility = "collapsed")

    # lob_row = st.columns([0.2, 0.7, 0.1])
    # with lob_row[0]:
    #     st.write("**LOB**")
    # with lob_row[1]:
    #     lob = st.selectbox('LOB',('Medicare', 'Medicaid'), label_visibility = "collapsed")

    field_row = st.columns([0.2, 0.7, 0.1])
    with field_row[0]:
        st.write("**Field Group**")
    with field_row[1]:
        field_group = st.selectbox('Field Group',('Unique Key', 'Contract Related', 'Pricing Before Carveouts - I'
            , 'Pricing Before Carveouts - II', 'Carveout Indicator, Code Type and Code #s - I'
            , 'Carveout Indicator, Code Type and Code #s - II', 'Carveout Indicator, Code Type and Code #s - III'
            , 'Optimize Carving Indic.', 'Carveout Method - I', 'Carveout Method - II', 'Provider'
            , 'Timeline'), label_visibility = "collapsed")

    # llm_row = st.columns([0.2, 0.7, 0.1])
    # with llm_row[0]:
    #     st.write("**Langauge Model**")
    # with llm_row[1]:
    #     llm_selected = st.selectbox('Langauge Model',('Claude 2', 'Claude Instant', 'Llama 2 Chat 70B'
    #         , 'Titan Text Express'), index=1, label_visibility = "collapsed")
    llm_selected = 'Claude Instant'
    
    if field_group == 'Unique Key':
        fields = fields[fields['PRIORITY'] == 'A']
    elif field_group == 'Contract Related':
        fields = fields[fields['PRIORITY'] == 'C']
    elif field_group == 'Pricing Before Carveouts - I':
        fields = fields[fields['PRIORITY'] == 'B']
        fields = np.array_split(fields, 2)[0]
    elif field_group == 'Pricing Before Carveouts - II':
        fields = fields[fields['PRIORITY'] == 'B']
        fields = np.array_split(fields, 2)[1]
    elif field_group == 'Carveout Indicator, Code Type and Code #s - I':
        fields = fields[fields['PRIORITY'] == 'F']
        fields = np.array_split(fields, 3)[0]
    elif field_group == 'Carveout Indicator, Code Type and Code #s - II':
        fields = fields[fields['PRIORITY'] == 'F']
        fields = np.array_split(fields, 3)[1]
    elif field_group == 'Carveout Indicator, Code Type and Code #s - III':
        fields = fields[fields['PRIORITY'] == 'F']
        fields = np.array_split(fields, 3)[2]
    elif field_group == 'Carveout Methodology - I':
        fields = fields[fields['PRIORITY'] == 'G']
        fields = np.array_split(fields, 4)[0]
    elif field_group == 'Carveout Methodology - II':
        fields = fields[fields['PRIORITY'] == 'G']
        fields = np.array_split(fields, 4)[1]
    elif field_group == 'Carveout Method - III':
        fields = fields[fields['PRIORITY'] == 'G']
        fields = np.array_split(fields, 4)[2]
    elif field_group == 'Carveout Method - IV':
        fields = fields[fields['PRIORITY'] == 'G']
        fields = np.array_split(fields, 4)[3]
    elif field_group == 'Provider':
        fields = fields[fields['PRIORITY'] == 'D']
    elif field_group == 'Timeline':
        fields = fields[fields['PRIORITY'] == 'E']
    
    fields['Interrogation Question?'] = fields['Interrogation Question?'].fillna(' ')
    field_prompt_mapping = dict(zip(fields['Field Name'], fields['Interrogation Question?']))

    # Setup bedrock
    bedrock_runtime = boto3.client(
        service_name="bedrock-runtime",
        region_name="us-east-1"
    )

    # question_list = list(field_prompt_mapping.items())
    # class OutputSchema(BaseModel):
    #     question_list[0]: str
    #     question_list[1]: str
    #     question_list[2]: str
    #     question_list[3]: str
    #     question_list[4]: str
    #     question_list[5]: str
    #     question_list[6]: str
    #     question_list[7]: str
    #     question_list[8]: str
    #     question_list[9]: str
    #     question_list[10]: str
    #     question_list[11]: str
    #     question_list[12]: str
    #     question_list[13]: str
    #     question_list[14]: str
    #     question_list[15]: str
    #     question_list[16]: str
    #     question_list[17]: str
    #     question_list[18]: str
    #     question_list[19]: str

    # class OutputList(BaseModel):
    #     answer: List[OutputSchema]

    # question = '\\n'.join(question_list)
    question = json.dumps(field_prompt_mapping)
    # question_with_schema = f'{question}{OutputList.schema_json()}'
    question_with_schema = question

    def run_llm(bucket, file_name, llm_selected, field_values):
        # with open(os.path.join(SOURCE_DIRECTORY, file_name), 'r') as infile:
        #     context = infile.read()
        # # page_count = context.count('Start of Page No. = ')
        data = s3_client.get_object(Bucket=bucket, Key=file_name)
        contents = data['Body'].read()
        context = contents.decode("utf-8")

        if llm_selected == "Titan Text Express":
            context = context[:16000]
            prompt_data = f"""Answer the question based only on the information provided between ## and give step by step guide. You must answer in correct JSON format.
            #
            {context}
            #

            Question: {question}
            Answer: Answer in JSON format: {{
            """
            parameters = {
                "maxTokenCount":1024,
                "stopSequences":[],
                "temperature":0,
                "topP":0.9
                }

            body = json.dumps({"inputText": prompt_data, "textGenerationConfig": parameters})
            model_id = "amazon.titan-text-express-v1"  # change this to use a different version from the model provider

        elif llm_selected == 'Llama 2 Chat 70B':
            context = context[:6000]
            prompt_data = f"""Answer the question based only on the information provided between ## and give step by step guide. You must answer in correct JSON format.
            ##
            {context}
            ##

            Question: {question}
            Answer: Answer in JSON format: {{
            """
            payload={
                "prompt":"[INST]"+ prompt_data +"[/INST]",
                "max_gen_len":1024,
                "temperature":0.0,
                "top_p":0.9
            }
            body=json.dumps(payload)
            model_id="meta.llama2-70b-chat-v1"

        elif llm_selected in ['Claude Instant', 'Claude 2']:
            if llm_selected == 'Claude Instant':
                context = context[:150000]

            prompt_data = f"""

            Human: Use the following pieces of context to provide a concise answer to the questions at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. You must answer in JSON format.

            {context}

            Question: {question_with_schema}

            Assistant: Answer in JSON format: {{
            """
            body = json.dumps(
                {"prompt": anthropic.HUMAN_PROMPT + prompt_data + anthropic.AI_PROMPT, 
                 "max_tokens_to_sample": 1024,
                 "temperature":0.0,
                 "top_p":1,
                 "top_k":250,
                 "stop_sequences":[anthropic.HUMAN_PROMPT]
                  })
            if llm_selected == "Claude 2":
                model_id = "anthropic.claude-v2:1"
            else:
                model_id = "anthropic.claude-instant-v1"


        # def claude_prompt_format(prompt: str) -> str:
        #     # Add headers to start and end of prompt
        #     return "\n\nHuman: " + prompt + "\n\nAssistant:"

        # # Call Claude model
        # def call_claude(prompt):
        #     prompt_config = {
        #         "prompt": claude_prompt_format(prompt),
        #         "max_tokens_to_sample": 4096,
        #         "temperature": 0.5,
        #         "top_k": 250,
        #         "top_p": 0.5,
        #         "stop_sequences": [],
        #     }

        #     body = json.dumps(prompt_config)

        #     modelId = "anthropic.claude-instant-v1"
        #     accept = "application/json"
        #     contentType = "application/json"

        #     response = bedrock_runtime.invoke_model(
        #         body=body, modelId=modelId, accept=accept, contentType=contentType
        #     )
        #     response_body = json.loads(response.get("body").read())

        #     results = response_body.get("completion")
        #     return results

        # prompt = SECOND_PROMPT
        # result = call_claude(prompt)
        # st.write(result)

        df = pd.DataFrame(columns=['Contract Name','Field Name','Snippet','Page Number',
            'Field Extracted Value','Imputed Value'])
        # field_list = list(field_prompt_mapping.keys())
        # query_list = [field_prompt_mapping[x] for x in field_list]
        # st.write(field_prompt_mapping)

        # st.write(prompt_data)

        try:
            response = bedrock_runtime.invoke_model(
                body=body, 
                modelId=model_id, 
                accept="application/json", 
                contentType="application/json"
            )
            response_body = json.loads(response.get("body").read())

            if llm_selected == "Titan Text Express":
                response_text = response_body.get("results")[0].get("outputText")
            elif llm_selected == 'Llama 2 Chat 70B':
                response_text = response_body['generation']
            elif llm_selected in ['Claude Instant', 'Claude 2']:
                response_text = response_body['completion']
                    
        except:
            response_text = "failed"

        raw_response_text = response_text           
        response_text = response_text.strip()

        try:
            if response_text.split("{",1)[1].strip()[0] == '"':
                response_text = "{" + response_text.split("{",1)[1]
            else:
                response_text = "{" + response_text
        except:
            response_text = "{" + response_text
        if len(response_text.split("}",1)) > 1:
            if response_text.rsplit("}",1)[0].strip()[-1] == '"':
                response_text = response_text.rsplit("}",1)[0] + "}"
        else:
            response_text = response_text.rstrip(",")
            response_text = response_text + "}"
        try:
            response_dict = json.loads(response_text)
        except:
            response_dict = {"Test value": "Failed to extract"}

        # st.write(response_dict)
        field_list = list(response_dict.keys())
        answer_list = list(response_dict.values())
        # st.write(answer_list)
        location_list = [context.find(answer) if isinstance(answer, str) and answer != "" else -1 for answer in answer_list]
        snippet_list = [' '.join(context[:location].split('.')[-4:]) + ' ' + ' '.join(context[location:].split('. ')[:5]
            ) if location != -1 else ' ' for location in location_list]
        page_no_list = [" " if location == -1 else context[:location].rsplit("Start of Page No. = ", 1)[1] if len(context[:location].rsplit(
            "Start of Page No. = ", 1)) > 1 else context[:location].rsplit("Start of Page No. = ", 1)[0] for location in location_list]
        # st.write(location_list)
        page_no_list = [re.search(r'\d+', page).group() if page != " " and re.search(r'\d+', page) is not None else "" for page in page_no_list]
        

        # response_list = [QA({"query":query}) for query in query_list]
        # answer_list = [response['result'] for response in response_list]
        # doc_list = [response['source_documents'] for response in response_list]
        # snippet_list = [str(doc[0].page_content) for doc in doc_list]
        # page_no_list = [int(str(doc[0].metadata["source"]).rsplit('_page')[1].replace('.txt',''))+1 for doc in doc_list]

        df['Field Name'] = field_list
        df['Contract Name'] = file_name
        df['Snippet'] = snippet_list
        df['Page Number'] = page_no_list
        # df['Confidence Level'] = ' '
        df['Field Extracted Value'] = answer_list
        df = pd.merge(df, fields[['Field Name', 'SF_DB_COL_NAME']], how ='left', on ='Field Name')

        document_name = [x for x in list(field_values['Document_Name']) if not pd.isna(x) and file_name.rsplit('/',1)[1].replace(' MU','').replace(
            '_MU','').replace('.txt','') in x][0]
        field_values = field_values[field_values['Document_Name'] == document_name].head(1).transpose().reset_index()
        field_values.columns = ['SF_DB_COL_NAME', 'Actual Value']
        df = pd.merge(df, field_values, how ='left', on ='SF_DB_COL_NAME')
        df = df[['Contract Name','Field Name', 'SF_DB_COL_NAME', 'Snippet','Page Number', 'Field Extracted Value', 'Actual Value','Imputed Value']]
        return df, raw_response_text

    raw_response_text = '{"Test value": "Failed to extract"}'

    if st.button("Show Results"):

        if file_name == 'All':
            df_1 = pd.DataFrame(columns=['Contract Name','Field Name', 'SF_DB_COL_NAME', 'Snippet','Page Number'
                , 'Field Extracted Value', 'Actual Value','Imputed Value'])
            for contract in contract_list:
                print(contract)
                df, raw_response_text = run_llm(bucket, contract, llm_selected, field_values)
                df_1 = pd.concat([df_1, df], ignore_index = True)
        else:
            df_1, raw_response_text = run_llm(bucket, file_name, llm_selected, field_values)
        df_1.to_csv('temp2.csv', index=False)

    df2 = pd.read_csv('temp2.csv')
    df2['Imputed Value'] = ''
    edited_df = st.data_editor(df2)

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

    csv = convert_df(edited_df)

    buttons = st.columns(3)
    with buttons[0]:
        # st.button("Save All Imputations") 
        st.download_button("Download Table", csv, "file.csv", "text/csv", key='download-csv')
    with buttons[1]:
        # st.download_button("Download Table", csv, "file.csv", "text/csv", key='download-csv')
        st.write("")
    with buttons[2]:
        st.button("Kickoff Database Integration") 

    add_vertical_space(20)
    st.write(raw_response_text)

else:
    st.write("Access Denied")