from typing import List
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables.base import RunnableSequence
from langchain_openai import OpenAI
from langchain.globals import set_llm_cache

from app.model.transaction import Transaction
from app.schema.index import IncomeStatementLLMResponse
from config.index import config as env

from langchain_core.output_parsers import PydanticOutputParser

set_llm_cache(None)

def income_statement_prompt () -> ChatPromptTemplate:
    context_str = """
        You are an accountant skilled at organizing transactions from multiple different bank
        accounts and credit card statements to prepare an income statement.

        Input data is in the below csv format:
        transaction_date, category, name_description, amount, type\n
        {input_data_csv}

        Your task is to prepare an income statement. The output should be in the following format: {format_instructions}

        """
    prompt = ChatPromptTemplate.from_template(context_str)
    return prompt

async def call_llm(inputData: List[Transaction]) -> str:
    input_data_csv = '\n'.join(str(x) for x in inputData)

    output_parser = PydanticOutputParser(pydantic_object=IncomeStatementLLMResponse)

    prompt = income_statement_prompt().partial(format_instructions=output_parser.get_format_instructions())

    llm = OpenAI(name='Income Statement Generation Bot',
                        api_key=env.OPENAI_API_KEY,
                        # cache=True,
                        temperature=0.7,
                        verbose=True)

    try:
        runnable_chain = RunnableSequence(prompt, llm, output_parser)
    except Exception as e:
        print(f"runnable_chain error: {str(e)}")
        raise e

    try:
        output_chunks = runnable_chain.invoke({"input_data_csv": input_data_csv})
        return output_chunks

    except Exception as e:
        print(f"runnable_chain.invoke error: {str(e)}")
        raise e