# planner.py
from langchain_ollama import ChatOllama
from langchain_core.messages import SystemMessage, HumanMessage
import json
from typing import List, Dict
from config import MODEL_NAME

class PlannerAgent:
    def __init__(self, llm=None):
        self.llm = llm or ChatOllama(
            model=MODEL_NAME, 
            temperature=0.0
        )

    def create_plan(self, task: str, tools_description: List[Dict]) -> List[Dict]:
        system_prompt = """Ты - планировщик. Твоя задача - разбить запрос пользователя на последовательность шагов.
Каждый шаг - это вызов одного из доступных инструментов с конкретными аргументами.
Верни JSON массив шагов, где каждый шаг содержит:
  - "tool": имя инструмента (строго из списка)
  - "args": словарь аргументов для этого инструмента
  - "description": краткое пояснение, зачем этот шаг

Инструменты и их аргументы:
{tools_description}

План должен быть линейным (шаг за шагом). Если нужны условия или циклы - разворачивай их в последовательность.
После выполнения всех шагов задача должна быть решена.
Верни ТОЛЬКО JSON массив, без лишнего текста."""

        user_prompt = f"Задача: {task}"
        tools_str = json.dumps(tools_description, indent=2)

        messages = [
            SystemMessage(content=system_prompt.format(tools_description=tools_str)),
            HumanMessage(content=user_prompt)
        ]
        response = self.llm.invoke(messages)
        # парсим JSON
        plan = json.loads(response.content)
        return plan