from langchain.agents import create_agent
from langchain.tools import tool
from langchain_ollama import OllamaLLM

# --- LLM --------------------------------------------------------------------
llm = OllamaLLM(
    model="llama3",
    temperature=0.7,
    timeout=60,
)

# --- Инструмент -------------------------------------------------------------
@tool
def check_wish(wish: str) -> str:
    """Инструмент для проверки желания на наличие подвоха"""
    genie_agent = create_agent(
        model=llm,
        tools=[],
        system_prompt="""
Ты - коварный джинн, который ищет подвох в любом желании.
Проанализируй желание человека и найди скрытую опасность, буквальное толкование, неожиданные последствия.
Если подвох найден - коротко предупреди о нем, например:
"хочу много денег - деньги будут фальшивыми."
"хочу деньги на счет в банке - хорошо, но банк завтра обанкротится"

Если желание безопасно и не имеет подвоха - ответь "Желание безопасно! Исполняю."
Ты должен говорить только на русском языке."""
    )

    result = genie_agent.invoke(
        {
            "messages": [
                {"role": "human", "content": f"Проверь желание: {wish}"}
            ]
        }
    )
    return result["messages"][-1].content


# --- Агент человека ---------------------------------------------------------
human_agent = create_agent(
    model=llm,
    tools=[check_wish],
    system_prompt="""
Ты - человек, который загадывает желания джинну.
Твоя задача - передать желание джинну через инструмент check_wish и сообщить результат.
Говори только на русском языке."""
)

# --- Ввод -------------------------------------------------------------------
current_wish = "Хочу читать мысли"

print("ДЖИНН ГОТОВ ИСПОЛНЯТЬ ЖЕЛАНИЯ!")
print(f"Человек: {current_wish}\n")
print("Джинн:", end=" ", flush=True)


# --- Форматирование ---------------------------------------------------------
def format_message(message) -> str:
    """Возвращает строку, которую нужно вывести в консоль."""
    if message.content:
        return message.content
    if message.tool_calls:
        tool_call = message.tool_calls[0]
        return f"{tool_call.name}({tool_call.args})"
    return ""


step = 1

def format_chunk_message(chunk):
    """Обрабатывает чанк типа 'messages'."""
    global step
    message, meta = chunk

    if meta["langgraph_step"] != step:
        step = meta["langgraph_step"]
        print("\n --- --- --- \n")

    if message.content:
        # выводим токен без перевода строки, чтобы текст «тёк»
        print(message.content, end="", flush=True)


# --- Потоковый вызов --------------------------------------------------------
stream = human_agent.stream(
    {
        "messages": [
            {
                "role": "human",
                "content": (
                    f"Вот мое желание: '{current_wish}'. Проверь его у джинна "
                    "через инструмент check_wish и скажи мне результат."
                ),
            }
        ]
    },
    stream_mode=["messages", "updates"],
)

full_response = ""

for chunk_type, chunk_data in stream:
    if chunk_type == "messages":
        format_chunk_message(chunk_data)
        message, _ = chunk_data
        full_response += message.content or ""
    elif chunk_type == "updates":
        # событие завершения шага (например, вызов инструмента)
        if chunk_data.get("model"):
            last_msg = chunk_data["model"]["messages"][-1]
            formatted = format_message(last_msg)
            if formatted:
                print(f"\n[Вызов инструмента: {formatted}]")
                full_response += f"\n[Вызов инструмента: {formatted}]\n"

print("\n")

# --- Итог -------------------------------------------------------------------
if "исполняю" in full_response.lower() or "безопасно" in full_response.lower():
    print(f" Желание исполнено! Финальная версия: {current_wish}")
else:
    print(f" Джинн отказался исполнять желание: {current_wish}")
    if full_response:
        clean_response = full_response.replace("\n", " ").strip()
        print(f" Причина: {clean_response}")
    else:
        print(" Джинн устал...")