#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Deep‑Search Agent from Scratch
• ищет в интернете (browser tool)
• генерирует виртуальные файлы (write_file tool)
• после завершения выгружает все файлы в реальную FS
"""

import os
import json
from pathlib import Path
from typing import Dict, List

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.tools import Tool
from langgraph.types import Command, interrupt
from tqdm import tqdm

# ────────────────────── 1. Конфиги ───────────────────────
OPENAI_MODEL = "gpt-4o-mini"
TEMPERATURE = 0.2
MAX_TOKENS = 800

# Папка для виртуальных файлов (внутри памяти)
VIRTUAL_FS_ROOT = Path("virtual_fs")


# ────────────────────── 2. Модели и промпты ───────────────────
class AgentState(dict):
    """Состояние агента – список сообщений + виртуальная FS."""

    messages: List[Dict[str, str]]   # chat‑history
    virtual_files: Dict[str, str]     # filename -> content


# --------------------- 2.1. Инструмент «browser» ---------------------
def browser_tool(query: str) -> str:
    """
    Поисковый инструмент.
    Реализован простым запросом к Bing Search API (или любому публичному).
    Для примера – возвращаем фиктивный ответ, но в реальном коде
    можно использовать `requests` + `bing-search-api`.
    """
    # Здесь просто пример. Замените на настоящий HTTP‑запрос.
    return f"Результаты поиска по запросу «{query}»:\n1) ...\n2) ..."


# --------------------- 2.2. Инструмент «write_file» ---------------------
def write_file_tool(filename: str, content: str) -> str:
    """
    Создаёт виртуальный файл (сохраняет в dict).
    Возвращает подтверждение.
    """
    AgentState.virtual_files[filename] = content
    return f"Файл '{filename}' создан/обновлён."


# --------------------- 2.3. Инструмент «finalise» ---------------------
def finalise_tool() -> str:
    """Вызывается в конце работы – выгружает виртуальную FS."""
    out_dir = Path("output_files")
    out_dir.mkdir(exist_ok=True)
    for fname, data in AgentState.virtual_files.items():
        p = out_dir / fname
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(data, encoding="utf-8")
    return f"Виртуальная FS выгружена в {out_dir}"


# ────────────────────── 3. Создание графа агента ───────────────────
def build_agent() -> StateGraph:
    llm = ChatOpenAI(temperature=TEMPERATURE, model=OPENAI_MODEL)

    # Промпт с инструкциями к инструментам
    prompt_template = """
Вы – исследовательский агент. Ваша задача: собрать информацию по заданной теме,
создать файлы с результатами и после завершения выгрузить их в файловую систему.

У вас есть три инструмента:
1️⃣ browser(query) — поиск в интернете.
2️⃣ write_file(filename, content) — создание виртуального файла.
3️⃣ finalise() — вызывается только один раз в конце работы и выгружает файлы.

Используйте формат JSON для вызова инструментов:

{
  "name": "<tool_name>",
  "arguments": { ... }
}

Пожалуйста, начните с выполнения необходимого поиска. После завершения
выполните finalise().

Тема: {topic}
"""

    prompt = ChatPromptTemplate.from_template(prompt_template)

    # Определяем инструменты как объекты LangGraph Tool
    browser = Tool(
        name="browser",
        func=browser_tool,
        description="Ищет информацию в интернете по запросу.",
        args_schema=None,  # простые строки
    )
    write_file = Tool(
        name="write_file",
        func=write_file_tool,
        description="Создаёт виртуальный файл с заданным содержимым.",
        args_schema=None,
    )
    finalise = Tool(
        name="finalise",
        func=finalise_tool,
        description="Выгружает все виртуальные файлы в реальную FS.",
        args_schema=None,
    )

    tools = [browser, write_file, finalise]

    # Структура состояния
    def start_state(topic: str) -> AgentState:
        return {
            "messages": [{"role": "system", "content": prompt_template}],
            "virtual_files": {},
        }

    # ── Узлы графа ────────────────────────────────────────
    @build_agent.node()
    def llm_node(state: AgentState) -> dict:
        """LLM генерирует сообщение (может быть вызовом инструмента)."""
        messages = state["messages"]
        response = llm.invoke(messages)
        # Добавляем ответ LLM в историю
        return {"messages": messages + [{"role": "assistant", "content": response.content}]}

    @build_agent.node()
    def tool_node(state: AgentState) -> dict:
        """Выполняет последний вызов инструмента."""
        last_msg = state["messages"][-1]
        # Предполагаем, что LLM вернул JSON‑объект с name/arguments
        import json as _json

        try:
            payload = _json.loads(last_msg["content"])
            tool_name = payload["name"]
            arguments = payload.get("arguments", {})
        except Exception:
            return {"messages": state["messages"]}  # ничего не делаем

        # Находим нужный инструмент
        tool = next(t for t in tools if t.name == tool_name)
        result = tool.invoke(arguments)

        # Добавляем результат в историю
        return {
            "messages": state["messages"]
            + [{"role": "tool", "content": result, "name": tool_name}]
        }

    # Создаём граф
    builder = StateGraph(AgentState)
    builder.add_node("llm", llm_node)
    builder.add_node("tool", tool_node)

    # Условия перехода: если сообщение содержит вызов инструмента → node «tool», иначе — «llm»
    def cond(state: AgentState):
        last_msg = state["messages"][-1]
        if last_msg.get("role") == "assistant" and "{ \"name\": \"" in last_msg[
            "content"
        ]:
            return "tool"
        else:
            return "llm"

    builder.set_entry_point("llm")
    builder.add_conditional_edges("llm", cond, {"tool": "tool", "llm": "llm"})
    builder.add_edge("tool", END)

    # Поставляем чекпоинтер (необязательно)
    memory = InMemorySaver()
    return builder.compile(checkpointer=memory)


# ────────────────────── 4. Запуск агента ───────────────────
def main():
    topic = input("\nВведите тему исследования: ").strip()
    agent_graph = build_agent()

    # Инициализируем потоковую генерацию
    config = {"configurable": {"thread_id": "search-agent-1"}}
    for chunk_type, chunk_data in agent_graph.stream(
        {"messages": [], "virtual_files": {}},  # стартовое состояние
        config=config,
        stream_mode=["messages"],
    ):
        if chunk_type == "messages":
            msg = chunk_data["content"]
            print(msg, end="", flush=True)

    print("\n\n=== Виртуальная файловая система: ===")
    for fname, content in agent_graph.get_state(config)["virtual_files"].items():
        print(f"  - {fname} ({len(content)} символов)")

    # После завершения графа – файлы уже выгружены командой `finalise()`


if __name__ == "__main__":
    main()