import os
from pdfminer import high_level
from langchain_astradb import AstraDBVectorStore
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAIEmbeddings
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI

ASTRA_DB_API_ENDPOINT = os.environ["ASTRA_DB_API_ENDPOINT"]
ASTRA_DB_APPLICATION_TOKEN = os.environ["ASTRA_DB_APPLICATION_TOKEN"]
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"]
GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]

collection_name = "ilj_test"

embedding = OpenAIEmbeddings(model="text-embedding-ada-002")

models = {
    "claude-3": ChatAnthropic(model='claude-3-sonnet-20240229'),
    "gemini-pro": ChatGoogleGenerativeAI(model="gemini-1.5-pro-latest", temperature=0.0)
}

def model_names():
    return models.keys()

def pipeline(file, model_name, balance_type, apsn_transactions, max_fees_per_day, min_overdrawn_fee, min_transaction_overdraft):
    disclosure_text = high_level.extract_text(file)
    
    # Read the legal cases context from the text file
    with open("all Text OCR.txt", "r") as f:
        legal_cases_context = f.read()

    astra = AstraDBVectorStore(
        api_endpoint=ASTRA_DB_API_ENDPOINT,
        token=ASTRA_DB_APPLICATION_TOKEN,
        collection_name=collection_name,
        embedding=embedding
    )

    related_docs = astra.search(disclosure_text, search_type="similarity")

    prompt = PromptTemplate.from_template(
        """
        law context:
        {context}

        above 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. Available balance or ledger balance should replace money in the first sentence.
        Here are the answers to the bank's policy questions:
        Do you charge on available balance or ledger balance?: {balance_type}
        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 disclosure text}}
        ------
        {{reasons_for_changes}}
        """
    )
    val = prompt.format(
        context=legal_cases_context,
        disclosure=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,
    )

    chat_response = models[model_name].invoke(input=val)

    return chat_response.content