import os

import cohere
import pandas as pd
import json
from unstructured.partition.pdf import partition_pdf

# uncomment to install dependencies
# !pip install cohere unstructured

# versions
print('cohere version:', cohere.__version__)

COHERE_API_KEY = os.environ.get("CO_API_KEY")
COHERE_MODEL = 'command-r-plus'
co = cohere.Client(api_key=COHERE_API_KEY)

def convert_to_json(text: str) -> dict:
    """
    Given text files, convert to json object and saves to csv.

    Args:
        text (str): The text to extract information from.

    Returns:
        dict: A dictionary containing the result of the conversion process.
    """

    MANDATORY_FIELDS = [
        "total_amount",
        "invoice_number",
    ]

    message = """# Instruction
    Given the text, convert to json object with the following keys:
    total_amount, invoice_number

    # Output format json:
    {{
        "total_amount": "<extracted invoice total amount>",
        "invoice_number": "<extracted invoice number>",
    }}

    Do not output code blocks.

    # Extracted PDF
    {text}
    """

    result = co.chat(
        message=message.format(text=text), model=COHERE_MODEL, preamble=None
    ).text

    try:
        result = json.loads(result)
        # check if all keys are present
        if not all(i in result.keys() for i in MANDATORY_FIELDS):
            return {"result": f"ERROR: Keys are missing. Please check your result {result}"}

        df = pd.DataFrame(result, index=[0])
        df.to_csv("output.csv", index=False)
        return {"result": "SUCCESS. All steps have been completed."}

    except Exception as e:
        return {"result": f"ERROR: Could not load the result as json. Please check the result: {result} and ERROR: {e}"}

def cohere_agent(
    message: str,
    preamble: str,
    verbose: bool = False,
) -> str:
    """
    Function to handle multi-step tool use api.

    Args:
        message (str): The message to send to the Cohere AI model.
        preamble (str): The preamble or context for the conversation.
        verbose (bool, optional): Whether to print verbose output. Defaults to False.

    Returns:
        str: The final response from the call.
    """

    functions_map = {
        "convert_to_json": convert_to_json,
    }

    tools = [
        {
            "name": "convert_to_json",
            "description": "Given a text, convert it to json object.",
            "parameter_definitions": {
                "text": {
                    "description": "text to be converted into json",
                    "type": "str",
                    "required": True,
                },
            },
        }
    ]

    counter = 1

    response = co.chat(
        model=COHERE_MODEL,
        message=message,
        preamble=preamble,
        tools=tools,
    )

    if verbose:
        print(f"\nrunning step 0")
        print(response.text)

    while response.tool_calls:
        tool_results = []

        if verbose:
            print(f"\nrunning step {counter}")
        for tool_call in response.tool_calls:
            print("tool_call.parameters:", tool_call.parameters)
            if tool_call.parameters:
                output = functions_map[tool_call.name](**tool_call.parameters)
            else:
                output = functions_map[tool_call.name]()

            outputs = [output]
            tool_results.append({"call": tool_call, "outputs": outputs})

            if verbose:
                print(
                    f"= running tool {tool_call.name}, with parameters: {tool_call.parameters}"
                )
                print(f"== tool results: {outputs}")

        response = co.chat(
            model=COHERE_MODEL,
            message="",
            chat_history=response.chat_history,
            preamble=preamble,
            tools=tools,
            tool_results=tool_results,
        )

        if verbose:
            print(response.text)
            counter += 1

    return response.text

def extract_pdf(path):
    """
    Function to extract text from a PDF file.
    """
    elements = partition_pdf(path)
    return "\n".join([str(el) for el in elements])


def pdf_extractor(pdf_path):
    """
    Main function that extracts pdf and calls the cohere agent.
    """
    pdf_text = extract_pdf(pdf_path)

    prompt = f"""
    # Instruction
    You are expert at extracting invoices from PDF. The text of the PDF file is given below.

    You must follow the steps below:
    1. Summarize the text and extract only the most information: total amount billed and invoice number.
    2. Using the summary above, call convert_to_json tool, which uses the summary from step 1.
    If you run into issues. Identifiy the issue and retry.
    You are not done unless you see SUCCESS in the tool output.

    # File Name:
    {pdf_path}

    # Extracted Text:
    {pdf_text}
    """
    output = cohere_agent(prompt, None, verbose=True)
    print(f"Finished extracting: {pdf_path}")

    print('Please check the output below')
    print(pd.read_csv('output.csv'))


pdf_extractor('simple_invoice.pdf')