from datetime import datetime

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.runnables.graph import CurveStyle
from langfuse.langchain import CallbackHandler
from langgraph.graph.state import CompiledStateGraph


def create_tool_message(messages, content="done", tool_call_index=0, tool_call_message_index=-1):
    tool_calls = messages[tool_call_message_index].tool_calls
    return {"role": "tool", "content": content, "tool_call_id": tool_calls[tool_call_index]['id']}


def create_ai_message(content: str | list[str | dict] | None = None):
    return AIMessage(content=content)


def create_user_message(content: str | list[str | dict] | None = None):
    return HumanMessage(content=content)


def create_system_message(content: str | list[str | dict] | None = None):
    return SystemMessage(content=content)


def dump_flow_result(result, user_content_block=False):
    if isinstance(result, dict):
        for m in result["messages"]:
            if user_content_block:
                for content_blocks in m.content_blocks:
                    print("---------------------------------------------------")
                    print(content_blocks)
            else:
                m.pretty_print()
    else:
        for e in result:
            if isinstance(e, tuple):
                e[0].pretty_print()
            else:
                messages = e.get("messages")
                if messages is not None:
                    messages[-1].pretty_print()


def create_run_config(run_name: str, thread_id_pattern="%Y%m%d-%H%M%S", recursion_limit=13, user_id="dev",
                      with_tracing=True):
    thread_id = datetime.now().strftime(thread_id_pattern)

    config = {
        "configurable": {"thread_id": thread_id},
        "recursion_limit": recursion_limit,
        "run_name": run_name,

    }

    if with_tracing:
        config["callbacks"] = [CallbackHandler()]
        config["metadata"] = {
            "langfuse_session_id": thread_id,
            "langfuse_user_id": user_id
        }

    return config


def save_mermaid(app: CompiledStateGraph, save_dir="./_data", xray=True, curve_style=CurveStyle.BASIS):
    with open(f"{save_dir}/{app.name}.mmd", "w", encoding="utf-8") as mmd_file:
        mmd_file.write(app.get_graph(xray=xray).draw_mermaid(curve_style=curve_style))
