from langchain.agents import create_agent
from langchain_community.chat_models import GigaChat
from langchain.tools import tool
import os

GPT2GIGA_PORT=8090
GIGACHAT_CREDENTIALS="MDE5YzBlMzUtNTJlYi03ODFiLTg1ZWUtZTc2MDFiZGUxYmM2OmNjZmIwODI2LTQ0ZTQtNDQwNC04NGE3LTgzNmM5ZDJhYmMzMg=="
GIGACHAT_SCOPE="GIGACHAT_API_B2B"
GIGACHAT_MODEL="GigaChat-MAX"
GIGACHAT_VERIFY_SSL_CERTS=False

llm = GigaChat(
    credentials=GIGACHAT_CREDENTIALS,
    scope=GIGACHAT_SCOPE,
    model=GIGACHAT_MODEL,
    verify_ssl_certs=False,  
    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.get('content'):
        return message['content']
    elif message.get('tool_calls'):
        tool_call = message['tool_calls'][0]
        return f"{tool_call['name']}({tool_call['args']})"
    return ""

step = 1

def format_chunk_message(chunk):
    global step
    message, meta = chunk
    
    if meta['langgraph_step'] != step:
        step = meta['langgraph_step']
        print('\n --- --- --- \n')
    
    if message.get('content'):
        print(message['content'], end='', flush=False)

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

full_response = ""
for chunk in stream:
    chunk_type, chunk_data = chunk
    
    if chunk_type == 'messages':
        format_chunk_message(chunk_data)
        message, _ = chunk_data
        if message.get('content'):
            full_response += message['content']
    
    elif chunk_type == 'updates':
        if chunk_data.get('model', None):
            last_message = chunk_data['model']['messages'][-1]
            formatted = format_message(last_message)
            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(" Джинн устал...")