# nodes.py

from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages
from langchain.tools import tool
from langgraph.prebuilt import create_react_agent, ToolNode
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.runnables import RunnableConfig

# from chatbot.llm_engine import llm_overall_agent, generate_expense_info_feedback
from llm_engine import (
    llm_overall_agent, 
    generate_expense_info_feedback, 
    generate_expense_description_feedback, 
    map_input_to_project_deliverable
)
from form_management.acceptance_criteria import acceptance_criteria
from form_management.project_task_descriptions import project_tasks
from prompts.agent_prompts import speaker_system_message
from form_management.form_management import Form
from kie import receipt_kie
import ast
import json
import sqlite3
import os
from collections import defaultdict

# conn = sqlite3.connect(r'C:\Users\Gavin\Desktop\hello-earth\db\project_data.db')
db_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'database', 'hello_earth_data_2.db')
conn = sqlite3.connect(db_path)
cursor = conn.cursor()

# placeholder, init empty form info
session_form_data = defaultdict(lambda: {
    "Seller Name": 'None',
    "Seller Address": 'None',
    "Seller Phone Number": 'None',
    "Buyer Name": 'None',
    "Buyer Address": 'None',
    "Transaction Date": 'None',
    "Total Payment Amount": 'None',
    "Associated Deliverable": 'None',
    "Expense Description": 'None',
})

# define the project_deliverables from the db
# cursor.execute("SELECT title FROM deliverables")
cursor.execute("SELECT title FROM Deliverable")
rows = cursor.fetchall()
deliverable_titles = [row[0] for row in rows]
print(f"deliverable_titles: {deliverable_titles}")

def get_form_info(session_id):
    return session_form_data[session_id]

def set_form_info(session_id, key, value):
    print("set form info")
    session_form_data[session_id][key] = value
    print(session_id)
    print(session_form_data[session_id])

class State(TypedDict):
    messages: Annotated[list, add_messages]
    session_id: str
    # form_info: dict

# @tool
# def placeholder_tool(input: str) -> str:
#     """
#     Placeholder tool that does nothing.
#     """
#     return "This is a placeholder tool."

# session_id: str might need to be added later when there are concurrent users
@tool
def apply_ocr(image_path: str, config: RunnableConfig) -> str:
    """
    Use this tool when the user provides an image path.
    Apply optical-character-recognition (OCR) to the image and initalize the form information.
    Feedback will then be provided on whether the form is complete or not along with the updated form info.
    """
    print("using apply_ocr tool")

    session_id = config["configurable"].get("thread_id")

    temp_form_info = receipt_kie(image_path)

    form_info = session_form_data[session_id]

    for key in temp_form_info.keys():
        if key in form_info:
            form_info[key] = temp_form_info[key]

    print(type(form_info))
    print(form_info)
 
    feedback = auditor_feedback(form_info)

    print("feedback: {} end of feedback".format(feedback))

    return f"""
    <form_status>
    {form_info}
    </form_status>

    <feedback>
    {feedback}
    </feedback>
    """

@tool
def edit_form(key: str, value: str, config: RunnableConfig) -> str:
    """
    This is used to edit the form data.
    The key must be one of the keys in the form_info dictionary.
    The target key will be assigned with the provided value.
    Feedback will then be provided on whether the form is complete or not along with the updated form info.
    Use this when the user wants to edit the form data.
    """
    print("using edit_form tool")

    session_id = config["configurable"].get("thread_id")

    form_info = session_form_data[session_id]

    if key in form_info:
        form_info[key] = value
        print(form_info)
        print(f"Updated {key} to {value}")
        if key == "Expense Description":
            feedback = expense_description_feedback(form_info, value)
        elif key == "Associated Deliverable":
            res = map_associated_deliverable(value, deliverable_titles)
            if res == "No Matching Deliverable":
                form_info[key] = "None"
                feedback = "A valid associated deliverable is required."
            else:
                form_info[key] = res
                feedback = auditor_feedback(form_info)
        else:
            feedback = auditor_feedback(form_info)
            
        print("feedback: {} end of feedback".format(feedback))
        return f"""
        <form_status>
        {form_info}
        </form_status>

        <feedback>
        {feedback}
        </feedback>
        """
    else:
        print(form_info)
        print(f"Invalid key: {key}. No changes made.")
        return f"Invalid key: {key}. No changes made."
    
@tool
def inspect_form(config: RunnableConfig) -> str:
    """
    This is used to inspect the form data.
    Use this when you need to address any inquiries the user might have about the current state of the form.
    """
    print("using inspect_form tool")

    session_id = config["configurable"].get("thread_id")
    form_info = session_form_data[session_id]
    print("session_id:",session_id)
    print("form_info:",form_info)

    return f"""
    <form_status>
    {form_info}
    </form_status>
    """

tools = [apply_ocr, edit_form, inspect_form]

speaker_memory = MemorySaver()
speaker_agent = create_react_agent(
    llm_overall_agent,
    tools,
    state_modifier=speaker_system_message,
    checkpointer=speaker_memory
)

tool_node = ToolNode(tools=tools)

def speaker_chatbot_node(state: State):
    config = {"configurable": {"thread_id": state["session_id"]}}

    response = speaker_agent.invoke(
        {
            "messages": [
                (
                    "user", "{}".format(state["messages"][-1].content)
                )
            ]
        },
        config,        
    )["messages"][-1]

    return {
        "messages": response.content,
        "session_id": state["session_id"]
    }

# provides feedback based on form info
def auditor_feedback(form_info):
    print("in auditor_feedback...")
    response = generate_expense_info_feedback(acceptance_criteria, form_info)
    return response

def expense_description_feedback(form_info, expense_description):
    response = generate_expense_description_feedback(form_info["Associated Deliverable"], expense_description)
    return response

# assume static project deliverables, would need project id to get the deliverables in actual case
def map_associated_deliverable(user_input, project_deliverables):
    response = map_input_to_project_deliverable(user_input, str(project_deliverables))
    return response

# # example usage
# if __name__ == "__main__":
#     # example_form_info = {
#     #     "Receipt Title": "Office Supplies",
#     #     "Project Name": "Project Alpha",
#     #     "Date": "2023-10-01",
#     #     "Receipt Number": "12345",
#     #     "Payer's Name": "John Doe",
#     #     "Payment Details Table": {
#     #         "Column Headers": ["Item", "Quantity", "Price"],
#     #         "Row Entries": [["Pen", "10", "$5"], ["Notebook", "5", "$20"]],
#     #     },
#     #     "Total Payment Amount": "$125",
#     #     "Incompleteness Description": "None",
#     #     "Expense Justification": "Office supplies for team use",
#     # }

#     # example_form_info = {
#     #     "Receipt Title": "Office Supplies",
#     #     "Project Name": "Reforestation of the Amazon Forest",
#     #     "Date": "2023-10-01",
#     #     "Receipt Number": "12345",
#     #     "Payer's Name": "John Doe",
#     #     "Payment Details Table": {
#     #         "Column Headers": ["Item", "Quantity", "Price"],
#     #         "Row Entries": [["Pen", "10", "$5"], ["Notebook", "5", "$20"]],
#     #     },
#     #     "Total Payment Amount": "$125",
#     #     "Incompleteness Description": "None",
#     #     "Expense Justification": "Office supplies for team use",
#     # }

#     # example_form_info = {
#     #     "Receipt Title": "Office Supplies",
#     #     "Project Name": "Reforestation of the Amazon Forest",
#     #     "Date": "2023-10-01",
#     #     "Receipt Number": "12345",
#     #     "Payer's Name": "John Doe",
#     #     "Payment Details Table": {
#     #         "Column Headers": ["Item", "Quantity", "Price"],
#     #         "Row Entries": [["Pen", "10", "$5"], ["Notebook", "5", "$20"]],
#     #     },
#     #     "Total Payment Amount": "$125",
#     #     "Incompleteness Description": "None",
#     #     "Expense Justification": "The office supplies are used for paperwork regarding the project.",
#     # }

#     # feedback = auditor_feedback(example_form_info)
#     # print("Generated Feedback:", feedback)