import os
import base64
import fitz  # PyMuPDF
from langchain_core.prompts import PromptTemplate
from google.cloud import aiplatform
from google.cloud.aiplatform_v1.services.model_service import ModelServiceClient
from google.cloud.aiplatform_v1.types import GenerateContentRequest, GenerationConfig
import streamlit as st

# Function to pad base64 strings
def pad_base64(base64_string):
    return base64_string + '=' * (-len(base64_string) % 4)

# Initialize the Google AI Platform
aiplatform.init(project="akroda", location="us-central1")

# Define the documents as dictionaries, ensuring correct padding
documents = [
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLjQKJeODgxNz5dL1Jvb3QgMTU0IDAgUi9TaXplIDE2Nj4+CnN0YXJ0eHJlZgoyMTY0NjkKJSVFT0YK"))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLjQKJeLjz9MKNijU+PgpzdGFydHhyZWYKMTMxMDY0CiUlRU9GCg=="))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLjQKJeLjz9MKNiAwZDU0YTVlNzllMWRhYWY1ZDQ2YjI+XS9Sb290IDE3NyAwIFIvU2l6ZSAxODc+PgpzdGFydHhyZWYKMjA3NTk5CiUlRU9GCg=="))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLjQKJeLjz9ML1Jvb3QgMTg5IDAgUi9TaXplIDE5OT4+CnN0YXJ0eHJlZgoxOTgzNzMKJSVFT0YK"))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xCcnCmVuZHN0cmVhbQplbmRvYmoKc3RhcnR4cmVmCjIwOTgyNQolJUVPRgo="))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLj+CnN0YXJ0eHJlZgoyMTk5MDYKJSVFT0YK"))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLjQKJiUlRU9GCg=="))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLjQKJe90IDMwOCAwIFIvU2l6ZSAzMTg+PgpzdGFydHhyZWYKMjcwNzU3CiUlRU9GCg=="))},
    {"content_type": "application/pdf", "data": base64.b64decode(pad_base64("JVBERi0xLjUNJeLjz9MNCjcgMCBvYmoNPDwvTGluZWFyaXplZCAxL0wgNjc1NzgvTyA5L0UgNjAyNDYvTiAxL1QgNjcyODcvSCBbIDQ4MyAxNTRdPj4NZW5kb2JxDSAgICAgICAgICAgICAgICAgICAgDQoyMiAwIG9iag08PC9EZWNvZGVQYXJtczw8L0NvbHVtbnMgNC9QcmVkaWN0b3IgMTI+Pi9GaWx0ZXIvRmxhdGVEZWNvZGUvSURbPDE3NzU4MkJFODc4MzRFQjNBOEM3RkIzQTgyRjFFMEFCPjw5MzI2Qjk4REM4NjQ2RTRCODI3MzZFQUEzOENEQjFBQj5dL0luZGV4WzcgMjhdL0luZm8gNiAwIFIvTGVuZ3RoIDgzL1ByZXYgNjcyODgvUm9vdCA4IDAgUi9TaXplIDM1L1R5cGUvWFPRg0K"))}
]

text1 = """
attached are several cases and a bank disclosure. Using the cases, please provide changes to the disclosure and keep as much formatting as possible and to ensure there are no legal contradictions between the content of the disclosure and the cases and please provide reasoning for each proposed change. Please also integrate the bank's policies into the disclosure. In the first sentence, please include a reference to the account agreement "for more information on overdrafts" and a placeholder for a URL. 
Here are the answers to the bank's policy questions:
Do you charge on available balance or ledger balance?: {balance_type} (which should replace money in the first sentence)
Do you charge for APSN transactions?: {apsn_transactions}
How many overdraft fees per day can be charged?: {max_fees_per_day}
What is the minimum amount overdrawn to incur a fee?: ${min_overdrawn_fee}
What is the minimum transaction amount to trigger an overdraft?: ${min_transaction_overdraft}

Please output in the following format:
{{entire updated disclosure text with changes bolded}}
------
{{reasons for each change listed and cases cited}}
"""

prompt = PromptTemplate(
    input_variables=["context", "disclosure", "balance_type", "apsn_transactions", "max_fees_per_day", "min_overdrawn_fee", "min_transaction_overdraft"],
    template=text1,
)

# Placeholder values for the variables used in prompt formatting
legal_cases_context = "Provide the legal context here..."
disclosure_text = "Include the initial disclosure text here..."
balance_type = "available balance"
apsn_transactions = "yes"
max_fees_per_day = 3
min_overdrawn_fee = 5
min_transaction_overdraft = 1

# Base64 encode the disclosure text
encoded_disclosure_text = base64.b64encode(disclosure_text.encode()).decode()

val = prompt.format(
    context=legal_cases_context,
    disclosure=encoded_disclosure_text,
    balance_type=balance_type,
    apsn_transactions=apsn_transactions,
    max_fees_per_day=max_fees_per_day,
    min_overdrawn_fee=min_overdrawn_fee,
    min_transaction_overdraft=min_transaction_overdraft,
)

generation_config = GenerationConfig(
    max_output_tokens=8192,
    temperature=1,
    top_p=0.95,
)

def generate(document_parts, prompt_text):
    model_service_client = ModelServiceClient()
    model_resource_name = model_service_client.model_path("akroda", "us-central1", "gemini-1.5-pro-001")
    response = model_service_client.generate_content(
        request=GenerateContentRequest(
            model=model_resource_name,
            documents=document_parts,
            prompt=prompt_text,
            generation_config=generation_config,
        )
    )
    return response.generated_text

def pipeline(file, model_name, balance_type, apsn_transactions, max_fees_per_day, min_overdrawn_fee, min_transaction_overdraft):
    response_text = generate(documents, val)
    return response_text

# Streamlit Interface
st.title("Bank Disclosure Update Pipeline")
st.write("Upload your document and provide the necessary details to update the bank disclosure.")

uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")

if uploaded_file is not None:
    file_details = {"FileName": uploaded_file.name, "FileType": uploaded_file.type}
    st.write(file_details)
    
    # Extract text from the uploaded PDF file
    pdf_document = fitz.open(stream=uploaded_file.read(), filetype="pdf")
    content = ""
    for page_num in range(len(pdf_document)):
        page = pdf_document.load_page(page_num)
        content += page.get_text()
    st.text(content)

    # Placeholder for user inputs
    balance_type = st.text_input("Balance Type", "available balance")
    apsn_transactions = st.text_input("APSN Transactions", "yes")
    max_fees_per_day = st.number_input("Max Fees Per Day", min_value=1, value=3)
    min_overdrawn_fee = st.number_input("Min Overdrawn Fee ($)", min_value=0, value=5)
    min_transaction_overdraft = st.number_input("Min Transaction Overdraft ($)", min_value=0, value=1)

    if st.button("Generate Updated Disclosure"):
        # Run the pipeline with the provided inputs
        result = pipeline(uploaded_file, "gemini-1.5-pro-001", balance_type, apsn_transactions, max_fees_per_day, min_overdrawn_fee, min_transaction_overdraft)
        st.write("Updated Disclosure:")
        st.text(result)