# solution.py
# Планирующий агент на основе LangGraph, LangChain и questionary
# Запуск: python solution.py

import os
from typing import TypedDict, List, Optional, Dict, Any
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver
from questionary import select, text

# ----------------------------------------------------------------------
# 1. Определяем состояние графа
# ----------------------------------------------------------------------
class AgentState(TypedDict):
    # Ввод задачи от пользователя
    task: str
    # Список шагов, которые LLM распланировал
    plan: List[str]
    # Текущий индекс шага
    step_index: int
    # Список результатов выполнения шагов
    step_results: List[str]
    # Итоговый ответ
    final_answer: Optional[str]
    # Выбранный вариант (если есть)
    choice: Optional[str]

# ----------------------------------------------------------------------
# 2. Инициализируем LLM
# ----------------------------------------------------------------------
# Ожидается, что переменная окружения OPENAI_API_KEY установлена
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)

# ----------------------------------------------------------------------
# 3. Определяем функции-узлы
# ----------------------------------------------------------------------
def plan(state: AgentState) -> Dict[str, Any]:
    """
    Первый узел: LLM разбивает задачу на шаги.
    """
    prompt = f"""
    Задача: {state['task']}

    Разбей её на последовательные шаги, которые можно выполнить один за другим.
    Выведите список шагов, каждый шаг в отдельной строке, без префиксов.
    """
    response = llm.invoke(prompt)
    steps = [s.strip() for s in response.content.splitlines() if s.strip()]
    return {"plan": steps, "step_index": 0, "step_results": []}

def execute_step(state: AgentState) -> Dict[str, Any]:
    """
    Выполняем текущий шаг из плана.
    """
    idx = state["step_index"]
    if idx >= len(state["plan"]):
        # План завершён
        return {"final_answer": "План выполнен."}
    current_step = state["plan"][idx]
    prompt = f"""
    Выполните следующий шаг задачи:

    Шаг {idx + 1}: {current_step}

    Предоставьте краткий результат выполнения этого шага.
    """
    response = llm.invoke(prompt)
    result = response.content.strip()
    new_results = state["step_results"] + [result]
    return {
        "step_results": new_results,
        "step_index": idx + 1
    }

def ask_choice(state: AgentState) -> Dict[str, Any]:
    """
    Если в результате выполнения шага есть варианты выбора,
    предлагаем пользователю выбрать один.
    """
    last_result = state["step_results"][-1]
    # Простейший парсинг: ищем строки вида "1) ... 2) ..."
    import re
    matches = re.findall(r'\d+\)\s*([^\d]+)', last_result)
    if not matches:
        # Нет вариантов, просто продолжаем
        return {}
    # Показываем варианты пользователю
    choice = select(
        "Выберите вариант:",
        choices=[f"{i+1}) {m.strip()}" for i, m in enumerate(matches)]
    ).ask()
    # Сохраняем выбранный вариант
    return {"choice": choice}

def finalize(state: AgentState) -> Dict[str, Any]:
    """
    После завершения всех шагов формируем итоговый ответ.
    """
    prompt = f"""
    Задача: {state['task']}

    Выполнены шаги:
    {chr(10).join(f"{i+1}. {r}" for i, r in enumerate(state['step_results']))}

    На основании результата сформируйте окончательный ответ на задачу.
    """
    response = llm.invoke(prompt)
    return {"final_answer": response.content.strip()}

# ----------------------------------------------------------------------
# 4. Строим граф
# ----------------------------------------------------------------------
builder = StateGraph(AgentState)

# Добавляем узлы
builder.add_node("plan", plan)
builder.add_node("execute_step", execute_step)
builder.add_node("ask_choice", ask_choice)
builder.add_node("finalize", finalize)

# Определяем переходы
builder.set_entry_point("plan")
builder.add_conditional_edges(
    "plan",
    lambda x: "execute_step",
    {"execute_step": "execute_step"}
)
builder.add_conditional_edges(
    "execute_step",
    lambda x: "finalize" if x["step_index"] >= len(x["plan"]) else "ask_choice",
    {"finalize": "finalize", "ask_choice": "ask_choice"}
)
builder.add_conditional_edges(
    "ask_choice",
    lambda x: "execute_step",
    {"execute_step": "execute_step"}
)
builder.add_edge("finalize", END)

# Сохраняем состояние в SQLite (чекпоинтер)
checkpoint = SqliteSaver.from_conn_str("sqlite:///langgraph_checkpoint.db")
graph = builder.compile(checkpointer=checkpoint)

# ----------------------------------------------------------------------
# 5. Запуск агента
# ----------------------------------------------------------------------
def main() -> None:
    # Получаем задачу от пользователя
    task = text("Введите задачу для агента:").ask()
    # Инициализируем состояние
    init_state: AgentState = {
        "task": task,
        "plan": [],
        "step_index": 0,
        "step_results": [],
        "final_answer": None,
        "choice": None,
    }
    # Запускаем граф
    for event in graph.stream(init_state):
        # event - это обновлённое состояние
        if event["final_answer"]:
            print("\n=== Итоговый ответ ===")
            print(event["final_answer"])
            break

if __name__ == "__main__":
    main()