"""State management for LangGraph (actor-first)."""

from __future__ import annotations

import json
import logging
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, TypeVar

from pydantic import BaseModel, ConfigDict, Field
from rx.core import Observable  # type: ignore[attr-defined]
from rx.subject import BehaviorSubject  # type: ignore[attr-defined]

T = TypeVar("T", bound="GraphState")


class StateUpdateMode(Enum):
    REPLACE = "replace"
    MERGE = "merge"
    APPEND = "append"


class StateSnapshot(BaseModel):
    state: dict[str, Any]
    timestamp: datetime
    node_id: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class GraphState(BaseModel):
    messages: list[dict[str, Any]] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)
    current_node: str | None = None
    execution_count: int = 0
    error: str | None = None

    model_config = ConfigDict(validate_assignment=True)

    def update(
        self, updates: dict[str, Any], mode: StateUpdateMode = StateUpdateMode.MERGE
    ) -> None:
        if mode == StateUpdateMode.REPLACE:
            for key, value in updates.items():
                if hasattr(self, key):
                    setattr(self, key, value)
        elif mode == StateUpdateMode.MERGE:
            for key, value in updates.items():
                if hasattr(self, key):
                    current = getattr(self, key)
                    if isinstance(current, dict) and isinstance(value, dict):
                        current.update(value)
                    elif isinstance(current, list) and isinstance(value, list):
                        current.extend(value)
                        if key == "messages" and len(current) > 50:
                            self.messages = current[-50:]
                    else:
                        setattr(self, key, value)
        elif mode == StateUpdateMode.APPEND:
            for key, value in updates.items():
                if hasattr(self, key):
                    current = getattr(self, key)
                    if isinstance(current, list):
                        if isinstance(value, list):
                            current.extend(value)
                        else:
                            current.append(value)
                        if key == "messages" and len(current) > 50:
                            self.messages = current[-50:]

    def to_dict(self) -> dict[str, Any]:
        return {
            "messages": self.messages,
            "metadata": self.metadata,
            "current_node": self.current_node,
            "execution_count": self.execution_count,
            "error": self.error,
        }

    @classmethod
    def from_dict(cls: type[T], data: dict[str, Any]) -> T:
        return cls(**data)


class StateManager:  # pylint: disable=too-many-instance-attributes
    def __init__(
        self,
        initial_state: GraphState | None = None,
        checkpoint_dir: Path | None = None,
        enable_time_travel: bool = False,
    ):
        self.logger = logging.getLogger(__name__)
        self.state = initial_state or GraphState()
        self.checkpoint_dir = checkpoint_dir
        self.enable_time_travel = enable_time_travel
        self.state_stream = BehaviorSubject(self.state)
        self.history: list[StateSnapshot] = []
        self.max_history_size = 100
        self.checkpoint_interval = 10
        self.update_count = 0
        self.is_closed = False
        if self.checkpoint_dir:
            self.checkpoint_dir.mkdir(parents=True, exist_ok=True)

    def get_state(self) -> GraphState:
        return self.state

    def update_state(
        self,
        updates: dict[str, Any],
        mode: StateUpdateMode = StateUpdateMode.MERGE,
        node_id: str | None = None,
    ) -> GraphState:
        if self.is_closed:
            raise RuntimeError("StateManager is closed")
        if self.enable_time_travel:
            snapshot = StateSnapshot(
                state=self.state.to_dict(), timestamp=datetime.now(), node_id=node_id
            )
            self.history.append(snapshot)
            if len(self.history) > self.max_history_size:
                self.history = self.history[-self.max_history_size :]

        self.state.update(updates, mode)
        self.state.execution_count += 1
        self.state_stream.on_next(self.state)

        self.update_count += 1
        if self.checkpoint_dir and self.update_count % self.checkpoint_interval == 0:
            self._save_checkpoint()
        return self.state

    def _save_checkpoint(self) -> None:
        if not self.checkpoint_dir:
            return
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        checkpoint_file = self.checkpoint_dir / f"checkpoint_{timestamp}.json"
        checkpoint_data = {
            "state": self.state.to_dict(),
            "timestamp": timestamp,
            "update_count": self.update_count,
        }
        checkpoint_file.write_text(
            json.dumps(checkpoint_data, indent=2), encoding="utf-8"
        )
        self.logger.debug("Saved checkpoint: %s", checkpoint_file)

    def load_checkpoint(self, checkpoint_file: Path) -> None:
        if self.is_closed:
            raise RuntimeError("StateManager is closed")
        checkpoint_data = json.loads(checkpoint_file.read_text(encoding="utf-8"))
        self.state = GraphState.from_dict(checkpoint_data["state"])
        self.update_count = checkpoint_data.get("update_count", 0)
        self.state_stream.on_next(self.state)
        self.logger.info("Loaded checkpoint: %s", checkpoint_file)

    def get_latest_checkpoint(self) -> Path | None:
        if not self.checkpoint_dir:
            return None
        checkpoints = list(self.checkpoint_dir.glob("checkpoint_*.json"))
        if not checkpoints:
            return None
        return max(checkpoints, key=lambda p: p.stat().st_mtime)

    def time_travel(self, steps_back: int = 1) -> GraphState | None:
        if self.is_closed:
            raise RuntimeError("StateManager is closed")
        if not self.enable_time_travel or not self.history:
            return None
        if steps_back >= len(self.history):
            steps_back = len(self.history) - 1
        snapshot = self.history[-(steps_back + 1)]
        self.state = GraphState.from_dict(snapshot.state)
        self.state_stream.on_next(self.state)
        return self.state

    def get_state_observable(self) -> Observable:
        return self.state_stream

    def clear_history(self) -> None:
        self.history.clear()

    def reset(self, initial_state: GraphState | None = None) -> None:
        if self.is_closed:
            raise RuntimeError("StateManager is closed")
        self.state = initial_state or GraphState()
        self.update_count = 0
        self.history.clear()
        self.state_stream.on_next(self.state)

    def close(self) -> None:
        """Mark this manager as closed and complete the state stream.

        After calling ``close()``, no further state updates should be
        performed.  Any underlying checkpoint resources are released.
        """
        if self.is_closed:
            return
        self.is_closed = True
        self.state_stream.on_completed()
        self.logger.info("StateManager closed")
