# agent.py
import argparse
from langchain_openai.chat_models import ChatOpenAI
from langchain.agents import initialize_agent
from langchain.tools import Tool
from config import config

from tools.gaia_tool import GaiaQATool
from tools.wiki_tool import WikiSearchTool
from tools.youtube_species_tool import YouTubeSpeciesTool
from tools.youtube_transcription_tool import YouTubeTranscriptionTool
from tools.excel_tool import ExcelTool
from tools.codegen_tool import CodeGenTool
from tools.web_search_tool import WebSearchTool
from tools.string_tool import StringTool
from tools.chess_tool import ChessTool
from tools.group_tool import GroupTool
from tools.grocery_tool import GroceryTool

def main() -> None:
    parser = argparse.ArgumentParser(description="Многофункциональный QA агент")
    parser.add_argument(
        "--input", required=True,
        help="Вопрос, URL или команда для инструмента"
    )
    args = parser.parse_args()
    user_input = args.input

    # Инициализируем LLM
    llm = ChatOpenAI(
        model_name=config.OPENAI_MODEL,
        openai_api_key=config.OPENAI_API_KEY,
        temperature=0.2
    )

    # Создаём и сохраняем инстансы инструментов
    gaia = GaiaQATool()
    wiki = WikiSearchTool()
    species = YouTubeSpeciesTool()
    transcript = YouTubeTranscriptionTool()
    excel = ExcelTool()
    codegen = CodeGenTool()
    web = WebSearchTool()
    string_tool = StringTool()
    chess = ChessTool()
    group = GroupTool()
    grocery = GroceryTool()

    # Собираем список Tool(...) из инстансов
    tools = [
        Tool(name=gaia.name, func=gaia._run, description=gaia.description),
        Tool(name=wiki.name, func=wiki._run, description=wiki.description),
        Tool(name=species.name, func=species._run, description=species.description),
        Tool(name=transcript.name, func=transcript._run, description=transcript.description),
        Tool(name=excel.name, func=excel._run, description=excel.description),
        Tool(name=codegen.name, func=codegen._run, description=codegen.description),
        Tool(name=web.name, func=web._run, description=web.description),
        Tool(name=string_tool.name, func=string_tool._run, description=string_tool.description),
        Tool(name=chess.name, func=chess._run, description=chess.description),
        Tool(name=group.name, func=group._run, description=group.description),
        Tool(name=grocery.name, func=grocery._run, description=grocery.description),
    ]

    # Инициализируем агента в режиме zero-shot-react-description
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent="zero-shot-react-description",
        verbose=True,
        max_iterations=30,
        early_stopping_method="generate",  # или "generate"
        handle_parsing_errors=True  # при ошибках парсинга пробует ещё раз
    )

    # Запускаем и получаем ответ
    answer = agent.invoke(user_input)
    print("\nОтвет:\n", answer)

if __name__ == "__main__":
    main()