from langgraph.graph import StateGraph, START, END
from langchain_ollama import ChatOllama
from langgraph.types import Send
from typing_extensions import TypedDict
from typing import Annotated
from pydantic import BaseModel, Field
import operator


class Section(BaseModel):
    num: int = Field(description="章节序号")
    name: str = Field(description="章节标题，采用“【第三章】 迷雾重现”的形式")
    description: str = Field(
        description="章节描述，最前面是章节标题，后面是章节内容。章节标题采用“【第三章】 迷雾重现”的形式，单独一行")


class Sections(BaseModel):
    sections: list[Section] = Field(description="文章的各个子章节")


class State(TypedDict):
    storyLine: str
    sections: list[Section]
    completedSections: Annotated[list, operator.add]
    novel: str


class WorkerState(TypedDict):
    section: Section
    completedSections: Annotated[list, operator.add]


llm = ChatOllama(model="qwen2.5:7b")


def getWholeStory(state):
    prompt = "你是一位著名小说家，正在创作一部精彩的侦探小说，首先给出故事梗概，1000字左右。"
    state["storyLine"] = llm.invoke(prompt).content
    print("storyLine  " + "*" * 80)
    print(state["storyLine"])
    return state


def orchestrate(state):
    planner = llm.with_structured_output(Sections)
    result = planner.invoke("""你是一位著名小说家，正在创作一部精彩的侦探小说，根据下面的【故事梗概】，将小说分成10个章节，并给出章节的剧情发展。1000字左右。
                            【故事梗概】
                            """ + state["storyLine"])
    state["sections"] = result.sections
    print("sections  " + "*" * 80)
    for section in state["sections"]:
        print(str(section.name + '        ' + section.description))
    return state

def buildGraph():
    graphBuilder = StateGraph(State)

    graphBuilder.add_node("getWholeStory", getWholeStory)
    graphBuilder.add_node("orchestrate", orchestrate)

    graphBuilder.add_edge(START, "getWholeStory")
    graphBuilder.add_edge("getWholeStory", "orchestrate")

    graph = graphBuilder.compile()
    return graph


if __name__ == "__main__":
    graph = buildGraph()

    #from tools import showGraph

    #showGraph.showGraphInCode(graph, "graph.jpg")

    state: State = {}
    result = graph.invoke(state)