from typing import TypedDict, List, Optional
from langchain_openai import ChatOpenAI
from pydantic import SecretStr
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver

# ---------- LLM ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b",
    base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
    api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
    temperature=0.2,
)

# ---------- State ----------
class PlanningState(TypedDict):
    task: str
    plan: List[str] | None
    current_step: int
    results: List[str]

# ---------- Nodes ----------
def planning(state: PlanningState) -> PlanningState:
    prompt = (
        f"Разбей задачу '{state['task']}' на 3–6 конкретных шагов. "
        "Ответ в виде JSON массива строк, например:\n"
        '["Шаг 1", "Шаг 2", ...]'
    )
    response = llm.invoke(prompt).content
    try:
        import json
        plan = json.loads(response)
        if not isinstance(plan, list):
            raise ValueError
    except Exception:
        # fallback: simple split by newlines or numbers
        plan = [line.strip() for line in response.splitlines() if line.strip()]
    return {
        **state,
        "plan": plan,
        "current_step": 0,
        "results": [],
    }

def execution(state: PlanningState) -> PlanningState:
    step_idx = state["current_step"]
    if state["plan"] is None or step_idx >= len(state["plan"]):
        return state
    step_text = state["plan"][step_idx]
    # Execute the step (here we just echo it; replace with real logic)
    result = f"[Шаг {step_idx + 1}] Выполнено: {step_text}"
    new_results = state["results"] + [result]
    return {
        **state,
        "current_step": step_idx + 1,
        "results": new_results,
    }

def should_continue(state: PlanningState) -> str:
    if state["plan"] is None or state["current_step"] >= len(state["plan"]):
        return "finish"
    return "execute"

# ---------- Graph ----------
graph = StateGraph(PlanningState)
graph.add_node("planning", planning)
graph.add_node("execution", execution)

# Connect planning to the first execution step
graph.add_edge("planning", "execution")

# Conditional loop: after each execution, decide whether to continue or finish
graph.add_conditional_edges(
    "execution",
    should_continue,
    {
        "execute": "execution",
        "finish": END,
    },
)

graph.set_entry_point("planning")
workflow = graph.compile(checkpointer=InMemorySaver())

# ---------- Demo ----------
if __name__ == "__main__":
    task_text = input("Задача: ").strip()
    if not task_text:
        print("Нет задачи.")
        exit(0)

    initial_state: PlanningState = {
        "task": task_text,
        "plan": None,
        "current_step": 0,
        "results": [],
    }

    result = workflow.invoke(initial_state)
    plan = result["plan"]
    results = result["results"]

    print("\nПлан:")
    for i, step in enumerate(plan, start=1):
        print(f"{i}. {step}")

    print("\nРезультаты выполнения:")
    for r in results:
        print(r)

    final_summary = "\n".join(results)
    print("\nИтог:\n", final_summary)