# solution.py

"""
Самописный поисковый агент на основе LangChain.
Собирает информацию по запросу, создаёт виртуальные файлы,
выводит их содержимое в консоль и сохраняет на диск.
"""

from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor
import questionary
from rich.console import Console

# ------------------------------------------------------------------
# 1. Инструмент поиска
# ------------------------------------------------------------------
@tool("search", "Поиск в интернете по запросу")
def search_tool(query: str) -> list[str]:
    """
    Возвращает список URL‑ов, найденных по запросу.
    Для демонстрации используется простая имитация поисковой выдачи.
    """
    # В реальном приложении здесь можно использовать API поисковика.
    # Здесь просто формируем фиктивные ссылки на основе запроса.
    return [
        f"https://example.com/search?q={query.replace(' ', '+')}",
        f"https://another-site.org/{query.replace(' ', '-')}",
    ]

# ------------------------------------------------------------------
# 2. Виртуальная файловая система
# ------------------------------------------------------------------
class VirtualFileSystem:
    """
    Хранит файлы в памяти и умеет сохранять их на диск.
    """

    def __init__(self):
        # словарь: имя файла -> содержимое
        self.files = {}

    def add_file(self, name: str, content: str) -> None:
        """Добавляет файл в память."""
        self.files[name] = content

    def list_files(self) -> list[str]:
        """Возвращает список имён файлов."""
        return list(self.files.keys())

    def get_content(self, name: str) -> str | None:
        """Получить содержимое файла по имени."""
        return self.files.get(name)

    def save_to_disk(self, path: str = "./output") -> None:
        """
        Сохраняет все файлы на диск в указанную папку.
        Если папка не существует – создаёт её.
        """
        import os

        os.makedirs(path, exist_ok=True)
        for name, content in self.files.items():
            file_path = os.path.join(path, name)
            with open(file_path, "w", encoding="utf-8") as f:
                f.write(content)

# ------------------------------------------------------------------
# 3. Агент
# ------------------------------------------------------------------
def create_agent() -> AgentExecutor:
    """
    Создаёт агента LangChain с LLM и инструментами.
    """
    llm = ChatOpenAI(temperature=0)  # можно заменить на другую модель

    tools = [search_tool]
    agent_executor = AgentExecutor.from_llm_and_tools(llm, tools)
    return agent_executor

# ------------------------------------------------------------------
# 4. Основной цикл
# ------------------------------------------------------------------
def main() -> None:
    console = Console()

    # Запрос пользователя
    query = questionary.text("Введите запрос для поиска:").ask()
    if not query:
        console.print("[red]Запрос не введён. Завершение работы.[/]")
        return

    # Создаём виртуальную файловую систему
    vfs = VirtualFileSystem()

    # Запускаем агента
    agent = create_agent()
    result = agent.invoke({"input": query})

    # Предположим, что агент вернёт список URL‑ов и контент для каждого файла.
    # Для демонстрации будем использовать результат инструмента поиска.
    urls: list[str] = search_tool(query)

    for i, url in enumerate(urls, start=1):
        file_name = f"result_{i}.txt"
        content = f"URL: {url}\n\n(Содержимое страницы не загружено в демо-версии.)"
        vfs.add_file(file_name, content)

    # Выводим содержимое файлов
    console.print("\n[bold green]Созданные файлы:[/]")
    for name in vfs.list_files():
        console.print(f"\n[underline]{name}[/]")
        console.print(vfs.get_content(name))

    # Сохраняем на диск
    output_dir = "./output"
    vfs.save_to_disk(output_dir)
    console.print(f"\n[bold blue]Файлы сохранены в папке: {output_dir}[/]")


if __name__ == "__main__":
    main()