import json
from pathlib import Path
from typing import Any, Dict, List

# ---------- FileCheckpointSaver ----------
from langgraph.checkpoint.base import BaseCheckpointSaver

class FileCheckpointSaver(BaseCheckpointSaver):
    """Сохраняет чекпоинт в JSON-файл."""

    def __init__(self, filepath: str = "./checkpoint.json"):
        self.filepath = Path(filepath)
        self.filepath.parent.mkdir(parents=True, exist_ok=True)

    def get(self, config: Dict[str, Any]) -> Dict[str, Any] | None:
        if self.filepath.exists():
            with open(self.filepath, "r", encoding="utf-8") as f:
                return json.load(f)
        return None

    def put(self, config: Dict[str, Any], checkpoint: Dict[str, Any]) -> Dict[str, Any]:
        with open(self.filepath, "w", encoding="utf-8") as f:
            json.dump(checkpoint, f, indent=2, ensure_ascii=False)
        return config

    def list(self, config: Dict[str, Any]) -> List[str]:
        if self.filepath.exists():
            return [self.filepath.name]
        return []

# ---------- ConversationMemory ----------
class ConversationMemory:
    """Управляет историей разговора с файловой persistence."""

    def __init__(self, filepath: str = "./memory.json"):
        self.filepath = Path(filepath)
        self.history: List[Dict[str, str]] = self._load()

    def _load(self) -> List[Dict[str, str]]:
        if self.filepath.exists():
            with open(self.filepath, "r", encoding="utf-8") as f:
                return json.load(f)
        return []

    def add(self, role: str, content: str) -> None:
        self.history.append({"role": role, "content": content})
        self._save()

    def get_history(self, limit: int = 10) -> List[Dict[str, str]]:
        return self.history[-limit:]

    def _save(self) -> None:
        self.filepath.parent.mkdir(parents=True, exist_ok=True)
        with open(self.filepath, "w", encoding="utf-8") as f:
            json.dump(self.history, f, indent=2, ensure_ascii=False)

    def clear(self) -> None:
        self.history = []
        self._save()

# ---------- LangGraph Agent ----------
from langgraph.graph import StateGraph, START, END
from langgraph.graph import add_messages
from typing import Annotated, TypedDict

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    memory_summary: str

# Simple LLM placeholder – replace with real LLM
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(temperature=0.7)


def agent_node(state: AgentState, llm) -> AgentState:
    response = llm.invoke(state["messages"])
    return {"messages": [response]}

# Build graph
builder = StateGraph(AgentState)
builder.add_node("agent", agent_node)
builder.set_entry_point("agent")
builder.add_edge(START, "agent")
builder.add_edge("agent", END)
agent = builder.compile()

# ---------- CLI ----------

def chat_loop(agent, memory: ConversationMemory):
    thread_id = "default"
    while True:
        user_input = input("\nВы: ")
        if user_input.lower() in ["exit", "quit"]:
            break
        memory.add("user", user_input)
        config = {"configurable": {"thread_id": thread_id}}
        result = agent.invoke({"messages": [{"role": "human", "content": user_input}]}, config=config)
        assistant_message = result["messages"][-1].content
        memory.add("assistant", assistant_message)
        print(f"\nАгент: {assistant_message}")

if __name__ == "__main__":
    checkpointer = FileCheckpointSaver("./checkpoint.json")
    memory = ConversationMemory("./memory.json")
    chat_loop(agent, memory)