import os
import json
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from agent_core import DeepAgent
from tools import web_search, create_virtual_file, list_virtual_files, export_files

# Load environment variables
load_dotenv()

# LLM initialization
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
    api_key=os.getenv("JOURNAL_MCP_PAT"),
    temperature=0.3,
)

# Helper to wrap the LLM callable expected by DeepAgent
class LLMWrapper:
    def __init__(self, llm):
        self.llm = llm
    def __call__(self, messages):
        # langchain returns a list of Message objects; convert to dict
        # For simplicity, we use the first assistant message content
        response = self.llm(messages)
        # The wrapper expects a dict with 'content'
        return {"content": response["content"]}

llm_wrapper = LLMWrapper(llm)

# Instantiate agent with tools
agent = DeepAgent(llm_wrapper, [web_search, create_virtual_file, list_virtual_files, export_files])

# Example tasks
TASKS = [
    {
        "description": "Найди информацию о LangGraph и создай файл summary.md",
        "query": "LangGraph python framework",
        "filename": "summary.md",
    },
    {
        "description": "Найди топ-5 Python библиотек для работы с LLM и создай файл llm_libs.md",
        "query": "top python libraries for llm",
        "filename": "llm_libs.md",
    },
    {
        "description": "Найди что такое ReAct агент и создай файл react_agent.md",
        "query": "ReAct agent definition",
        "filename": "react_agent.md",
    },
]

for task in TASKS:
    print(f"\n=== {task['description']} ===")
    # Step 1: search
    search_result = agent.run(task["query"])
    # Step 2: create file with search result
    create_msg = agent.run(f"create_virtual_file {task['filename']} | {search_result}")
    print(create_msg)

# Export all virtual files to disk
export_msg = agent.run("export_files output")
print(export_msg)

if __name__ == "__main__":
    # The script already executed tasks above; nothing else needed.
    pass