import os
from typing import List, Dict, Any

# Библиотеки, которые должны быть установлены в окружении:
# - langchain
# - fastapi (если понадобится)
# - pydantic
# - rich (для красивого вывода)

try:
    from langchain.llms import OpenAI
    from langchain.prompts import PromptTemplate
    from langchain.chains import LLMChain
except ImportError as exc:
    raise RuntimeError(
        "Не удалось импортировать необходимые библиотеки. "
        "Убедитесь, что установлены langchain и его зависимости."
    ) from exc

# Конфигурация модели (можно заменить на другую)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError(
        "Переменная окружения OPENAI_API_KEY не задана. "
        "Пожалуйста, установите ключ API для OpenAI."
    )

# Основной класс планирующего агента
class PlanningAgent:
    """
    Простая реализация планировщика задач на основе LLM.
    Принимает список целей и генерирует последовательность действий,
    которые необходимо выполнить для достижения каждой цели.
    """

    def __init__(self, llm_model: str = "gpt-3.5-turbo"):
        self.llm = OpenAI(model=llm_model, temperature=0.7)
        # Шаблон запроса к LLM
        self.prompt_template = PromptTemplate(
            input_variables=["goal"],
            template=(
                "Ты планировщик задач. Для цели: '{goal}' "
                "предложи пошаговый план действий в виде списка.\n"
                "Ответ должен быть коротким и конкретным, каждый шаг на одной строке."
            ),
        )
        self.chain = LLMChain(llm=self.llm, prompt=self.prompt_template)

    def plan_for_goal(self, goal: str) -> List[str]:
        """
        Генерирует план действий для одной цели.
        Возвращает список строк – шагов плана.
        """
        raw_output = self.chain.run({"goal": goal})
        # Разделяем вывод по строкам и убираем пустые
        steps = [step.strip() for step in raw_output.splitlines() if step.strip()]
        return steps

    def plan_for_goals(self, goals: List[str]) -> Dict[str, List[str]]:
        """
        Генерирует планы для списка целей.
        Возвращает словарь: цель → список шагов.
        """
        plans = {}
        for goal in goals:
            plans[goal] = self.plan_for_goal(goal)
        return plans

# Тестовый пример использования
if __name__ == "__main__":
    agent = PlanningAgent()
    sample_goals = [
        "Запустить новый веб-сайт",
        "Оптимизировать производительность базы данных",
        "Подготовить презентацию для инвесторов"
    ]
    plans = agent.plan_for_goals(sample_goals)
    for goal, steps in plans.items():
        print(f"\nЦель: {goal}")
        for i, step in enumerate(steps, 1):
            print(f"  Шаг {i}: {step}")