from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langfuse.langchain import CallbackHandler

from .config import REGION
from .dynamo import langfuse, get_contexto
from .agent_bedrock import create_agent
from .tools import ReportTools


def main(user_query, history, model, base):
    """Main execution function."""
    contexto_data = get_contexto(base)

    id_mapping = {
        f"id_{i}": real_id
        for i, real_id in enumerate(contexto_data["items_disponiveis"], 1)
    }
    local_items = {
        local_id: contexto_data["items_disponiveis"][real_id]
        for local_id, real_id in id_mapping.items()
    }
    report_tools = ReportTools(id_mapping)

    if contexto_data["filter"] == "period":
        CONSULT_RULES = """To use the tools you must give the id of the correspondant data, which can be associated to a given month and year in the following format year_month, which:
        -Year is the year in 4 digits (2025,2024,2023,2022,2021,...)
        -Month is th two digit representation: 01,02,03,04,05,06,07,08,09,10,11,12
        The format of the dict is: {id1:year_month1,id2:year_month2...}
        Choose the correct id based on the following dict:
"""
    elif contexto_data["filter"] == "event":
        CONSULT_RULES = """To use the tools you must give the id of the correspondant data, which can be associated to a event, which is in the format "Name - City DD/MM/YYYY", where the last is a date in the format day/month/year. Theformat of elements in dict is {id1:event_description1,id2:event_description2...}"""
    else:
        CONSULT_RULES = """Wrong filter value, you must terminate the workflow and ask the user to contact the technical team"""

    SYSTEM_PROMPT = """ You are a analitical agent in Brazilian Portuguese, with acess to monthly reports about a specific company, specified in the context. You have access to tools that lets you consult present variables in table, you always have access to "context", which keeps inside answers to different questions, that you may consult as you desire.
    Do not access other variables besides the ones reported by the tool and "context".
    You currently  have access to data in a period specified in the context,  so only answer questions inside the time window.
    <context>
    """ + contexto_data["contexto"] + """
    </context>
    """ + CONSULT_RULES + """
    <correlation>
    """ + str(local_items) + """
    </correlation>
    Here is the chat history:""" + history + """
    Inside the "NPS" in data is some useful values to calculate the NPS, which includes "distribuicao".
    Inside of it are grades and the amount of people who given that grade.
    Grades from 0 to 6 are detractors.
    Grades from 7 to 8 are neutral.
    Grades from 9 to 10 are promoters.
    Calculate the percentage of them when prompted about NPS and then  calculate the nps using the following formula: NPS = %promoter - %detractor, never use the medium of the notes.

    You have access to the tools:
        -get_consolidated_keys: Given a id returns the column names inside of a entity of a given table element.
        - get_monthly_report: given a id and a variable name, either one listed in the previous tool output or "context", returns its value. Using "context" gives you a summarization of many answers of questions asked to the customers.
    Answer, in Brazilian Portuguese, to the user the best you can with the given information, if you don't know the answer or how to answer say so, only answer from what you know.
    Always consult the most recent information when a date is not given, like questions "Quanto é meu nps?" """

    langfuse_handler = CallbackHandler()
    agent = create_agent(model, REGION, tools=report_tools.as_tools())

    initial_state = {
        "messages": [
            SystemMessage(content=SYSTEM_PROMPT),
            HumanMessage(content=user_query),
        ],
        "current_step": "init",
    }

    config = {"callbacks": [langfuse_handler], "tags": [base]}
    final_state = agent.invoke(initial_state, config=config)

    total_input_tokens = 0
    total_output_tokens = 0
    for msg in final_state["messages"]:
        if isinstance(msg, AIMessage) and hasattr(msg, "usage_metadata") and msg.usage_metadata:
            total_input_tokens += msg.usage_metadata.get("input_tokens", 0)
            total_output_tokens += msg.usage_metadata.get("output_tokens", 0)

    langfuse.flush()
    return {
        "response": final_state["messages"][-1].content,
        "input_tokens": total_input_tokens,
        "output_tokens": total_output_tokens,
        "total_tokens": total_input_tokens + total_output_tokens,
    }


if __name__ == "__main__":
    main(
        "Liste o nps mês a mês desde maio 2025 até dezembro 2025",
        "",
        "anthropic.claude-sonnet-4-5-20250929-v1:0",
        "bacio_transacional_loja_app",
    )