"""StateStore — persistence layer for agent/pipeline state."""
from __future__ import annotations

from dataclasses import asdict, is_dataclass
from datetime import UTC, datetime
import hashlib
import json
from pathlib import Path
import sqlite3
from typing import Any

from logos_adk.config import StateConfig
from logos_adk.errors import ConfigError, StateBackendError, StateError
from logos_adk.persistence import (
    get_postgresql_schema_version,
    get_sqlite_schema_version,
    set_postgresql_schema_version,
    set_sqlite_schema_version,
    sqlite_table_columns,
)
from logos_adk.postgres_utils import pooled_connect, resolve_database_url, safe_identifier
from logos_adk.state import AgentState
from logos_adk.types import Message, ToolCall, ToolResult, normalize_content_parts


STATE_SCHEMA_KIND = "state_store"
STATE_SCHEMA_VERSION = 1


def _json_default(value: Any) -> Any:
    if isinstance(value, datetime):
        return value.isoformat()
    if is_dataclass(value):
        return asdict(value)
    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")


def _serialize_state_data(state: AgentState) -> dict[str, Any]:
    return {
        "messages": [
            {
                "role": message.role,
                "content": message.content,
                "content_parts": message.content_parts,
                "sender": message.sender,
                "tool_calls": [
                    {
                        "id": tool_call.id,
                        "name": tool_call.name,
                        "arguments": tool_call.arguments,
                    }
                    for tool_call in message.tool_calls
                ],
                "tool_result": (
                    {
                        "call_id": message.tool_result.call_id,
                        "content": message.tool_result.content,
                        "content_parts": message.tool_result.content_parts,
                        "error": message.tool_result.error,
                        "duration_ms": message.tool_result.duration_ms,
                        "status": message.tool_result.status,
                    }
                    if message.tool_result is not None
                    else None
                ),
            }
            for message in state.messages
        ],
        "session_id": state.session_id,
        "metadata": state.metadata,
        "version": state.version,
        "updated_at": state.updated_at.isoformat() if state.updated_at else None,
    }


def _deserialize_state_data(payload: dict[str, Any]) -> AgentState:
    messages = []
    for raw_message in payload.get("messages", []):
        tool_calls = [
            ToolCall(
                id=item["id"],
                name=item["name"],
                arguments=item.get("arguments", {}),
            )
            for item in raw_message.get("tool_calls", [])
        ]
        raw_tool_result = raw_message.get("tool_result")
        tool_result = None
        if raw_tool_result is not None:
            tool_result = ToolResult(
                call_id=raw_tool_result["call_id"],
                content=raw_tool_result.get("content"),
                content_parts=normalize_content_parts(raw_tool_result.get("content_parts")),
                error=raw_tool_result.get("error"),
                duration_ms=raw_tool_result.get("duration_ms", 0),
                status=raw_tool_result.get("status", "completed"),
            )
        messages.append(
            Message(
                role=raw_message["role"],
                content=raw_message.get("content", ""),
                content_parts=normalize_content_parts(raw_message.get("content_parts")),
                sender=raw_message.get("sender"),
                tool_calls=tool_calls,
                tool_result=tool_result,
            )
        )

    updated_at = payload.get("updated_at")
    return AgentState(
        messages=messages,
        session_id=payload.get("session_id"),
        metadata=payload.get("metadata", {}),
        version=int(payload.get("version", 0)),
        updated_at=datetime.fromisoformat(updated_at) if updated_at else None,
    )


def _clone_state(state: AgentState) -> AgentState:
    return _deserialize_state_data(
        json.loads(json.dumps(_serialize_state_data(state), default=_json_default))
    )


class StateStore:
    """Abstract base for state persistence."""

    def save(self, namespace: str, state: AgentState) -> AgentState:
        raise NotImplementedError

    def load(self, namespace: str, session_id: str | None) -> AgentState | None:
        raise NotImplementedError

    def history(self, namespace: str, session_id: str | None, *, limit: int = 50) -> list[dict[str, Any]]:
        raise NotImplementedError

    def export_state(
        self,
        namespace: str,
        session_id: str | None,
        *,
        version: int | None = None,
    ) -> dict[str, Any] | None:
        raise NotImplementedError

    def import_state(
        self,
        namespace: str,
        payload: dict[str, Any],
        *,
        session_id: str | None = None,
    ) -> AgentState:
        raise NotImplementedError

    def rollback(self, namespace: str, session_id: str | None, *, version: int) -> AgentState:
        raise NotImplementedError

    def describe(self) -> dict[str, Any]:
        raise NotImplementedError


class InMemoryStateStore(StateStore):
    """In-memory state store for testing and development."""

    def __init__(self) -> None:
        self._latest: dict[tuple[str, str], AgentState] = {}
        self._history: dict[tuple[str, str], list[AgentState]] = {}

    def _key(self, namespace: str, session_id: str | None) -> tuple[str, str]:
        return (namespace, session_id or "")

    def save(self, namespace: str, state: AgentState) -> AgentState:
        key = self._key(namespace, state.session_id)
        stored_state = _clone_state(state)
        stored_state.version = len(self._history.get(key, [])) + 1
        stored_state.updated_at = datetime.now(UTC)
        self._latest[key] = stored_state
        self._history.setdefault(key, []).append(_clone_state(stored_state))
        return _clone_state(stored_state)

    def load(self, namespace: str, session_id: str | None) -> AgentState | None:
        state = self._latest.get(self._key(namespace, session_id))
        return _clone_state(state) if state is not None else None

    def history(self, namespace: str, session_id: str | None, *, limit: int = 50) -> list[dict[str, Any]]:
        versions = list(reversed(self._history.get(self._key(namespace, session_id), [])))[:limit]
        return [
            {
                "version": item.version,
                "saved_at": item.updated_at.isoformat() if item.updated_at else "",
                "session_id": item.session_id,
            }
            for item in versions
        ]

    def export_state(
        self,
        namespace: str,
        session_id: str | None,
        *,
        version: int | None = None,
    ) -> dict[str, Any] | None:
        versions = self._history.get(self._key(namespace, session_id), [])
        if not versions:
            return None
        if version is None:
            state = versions[-1]
        else:
            state = next((item for item in versions if item.version == version), None)
            if state is None:
                return None
        return {
            "agent_name": namespace,
            "session_id": state.session_id,
            "version": state.version,
            "saved_at": state.updated_at.isoformat() if state.updated_at else "",
            "state": _serialize_state_data(state),
        }

    def import_state(
        self,
        namespace: str,
        payload: dict[str, Any],
        *,
        session_id: str | None = None,
    ) -> AgentState:
        state_payload = payload.get("state", payload)
        state = _deserialize_state_data(state_payload)
        state.session_id = session_id if session_id is not None else state.session_id
        return self.save(namespace, state)

    def rollback(self, namespace: str, session_id: str | None, *, version: int) -> AgentState:
        versions = self._history.get(self._key(namespace, session_id), [])
        restored = next((item for item in versions if item.version == version), None)
        if restored is None:
            raise StateError(f"No saved state version {version} for session {session_id}")
        rolled_back = _clone_state(restored)
        rolled_back.session_id = session_id
        return self.save(namespace, rolled_back)

    def describe(self) -> dict[str, Any]:
        return {"backend": "in_memory", "shards": 1}


class SQLiteStateStore(StateStore):
    """Persistent state store backed by SQLite."""

    def __init__(self, path: str, table_name: str = "agent_state_versions") -> None:
        self.path = path
        self.table_name = safe_identifier(table_name)
        self._schema_ready = False

    def _connect(self) -> sqlite3.Connection:
        try:
            Path(self.path).parent.mkdir(parents=True, exist_ok=True)
            connection = sqlite3.connect(self.path)
            connection.row_factory = sqlite3.Row
            return connection
        except Exception as exc:
            raise StateBackendError(f"Failed to open SQLite state backend: {exc}") from exc

    def _ensure_schema(self, connection: sqlite3.Connection) -> None:
        if self._schema_ready:
            return
        connection.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                namespace TEXT NOT NULL,
                session_id TEXT NOT NULL,
                version INTEGER NOT NULL,
                saved_at TEXT NOT NULL,
                state_json TEXT NOT NULL,
                PRIMARY KEY (namespace, session_id, version)
            )
            """
        )
        columns = self._get_columns(connection)
        if {"agent_name", "session_key", "session_id", "version", "saved_at", "state_json"}.issubset(columns):
            self._migrate_legacy_schema(connection)
            columns = self._get_columns(connection)
        if {"namespace", "session_id", "version", "saved_at", "state_json"}.issubset(columns):
            connection.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_latest
                ON {self.table_name} (namespace, session_id, version DESC)
                """
            )
        recorded_version = get_sqlite_schema_version(connection, STATE_SCHEMA_KIND)
        if recorded_version != STATE_SCHEMA_VERSION:
            set_sqlite_schema_version(connection, STATE_SCHEMA_KIND, STATE_SCHEMA_VERSION)
        connection.commit()
        self._schema_ready = True

    def _session_key(self, session_id: str | None) -> str:
        return session_id or ""

    def _get_columns(self, connection: sqlite3.Connection) -> set[str]:
        return sqlite_table_columns(connection, self.table_name)

    def _uses_legacy_schema(self, connection: sqlite3.Connection) -> bool:
        columns = self._get_columns(connection)
        return "namespace" not in columns and "agent_name" in columns and "session_key" in columns

    def _migrate_legacy_schema(self, connection: sqlite3.Connection) -> None:
        legacy_table = f"{self.table_name}__legacy_v0"
        connection.execute(f"ALTER TABLE {self.table_name} RENAME TO {legacy_table}")
        connection.execute(
            f"""
            CREATE TABLE {self.table_name} (
                namespace TEXT NOT NULL,
                session_id TEXT NOT NULL,
                version INTEGER NOT NULL,
                saved_at TEXT NOT NULL,
                state_json TEXT NOT NULL,
                PRIMARY KEY (namespace, session_id, version)
            )
            """
        )
        connection.execute(
            f"""
            INSERT INTO {self.table_name} (namespace, session_id, version, saved_at, state_json)
            SELECT agent_name, session_key, version, saved_at, state_json
            FROM {legacy_table}
            """
        )
        connection.execute(f"DROP TABLE {legacy_table}")

    def save(self, namespace: str, state: AgentState) -> AgentState:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                session_id = self._session_key(state.session_id)
                current_version = connection.execute(
                    f"""
                    SELECT COALESCE(MAX(version), 0)
                    FROM {self.table_name}
                    WHERE namespace = ? AND session_id = ?
                    """,
                    (namespace, session_id),
                ).fetchone()[0]
                stored_state = _clone_state(state)
                stored_state.version = int(current_version) + 1
                stored_state.updated_at = datetime.now(UTC)
                payload = json.dumps(_serialize_state_data(stored_state), default=_json_default)
                connection.execute(
                    f"""
                    INSERT INTO {self.table_name} (namespace, session_id, version, saved_at, state_json)
                    VALUES (?, ?, ?, ?, ?)
                    """,
                    (
                        namespace,
                        session_id,
                        stored_state.version,
                        stored_state.updated_at.isoformat(),
                        payload,
                    ),
                )
                connection.commit()
                return _clone_state(stored_state)
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to persist state: {exc}") from exc

    def _fetch_state_row(
        self,
        namespace: str,
        session_id: str | None,
        *,
        version: int | None = None,
    ) -> sqlite3.Row | None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            session_key = self._session_key(session_id)
            if version is None:
                return connection.execute(
                    f"""
                    SELECT version, saved_at, state_json
                    FROM {self.table_name}
                    WHERE namespace = ? AND session_id = ?
                    ORDER BY version DESC
                    LIMIT 1
                    """,
                    (namespace, session_key),
                ).fetchone()
            return connection.execute(
                f"""
                SELECT version, saved_at, state_json
                FROM {self.table_name}
                WHERE namespace = ? AND session_id = ? AND version = ?
                LIMIT 1
                """,
                (namespace, session_key, version),
            ).fetchone()

    def load(self, namespace: str, session_id: str | None) -> AgentState | None:
        try:
            row = self._fetch_state_row(namespace, session_id)
            if row is None:
                return None
            state = _deserialize_state_data(json.loads(row["state_json"]))
            state.version = int(row["version"])
            state.updated_at = datetime.fromisoformat(row["saved_at"])
            return state
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to load state: {exc}") from exc

    def history(self, namespace: str, session_id: str | None, *, limit: int = 50) -> list[dict[str, Any]]:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                rows = connection.execute(
                    f"""
                    SELECT version, saved_at
                    FROM {self.table_name}
                    WHERE namespace = ? AND session_id = ?
                    ORDER BY version DESC
                    LIMIT ?
                    """,
                    (namespace, self._session_key(session_id), limit),
                ).fetchall()
            return [
                {
                    "version": int(row["version"]),
                    "saved_at": row["saved_at"],
                    "session_id": session_id,
                }
                for row in rows
            ]
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to list state history: {exc}") from exc

    def export_state(
        self,
        namespace: str,
        session_id: str | None,
        *,
        version: int | None = None,
    ) -> dict[str, Any] | None:
        try:
            row = self._fetch_state_row(namespace, session_id, version=version)
            if row is None:
                return None
            return {
                "agent_name": namespace,
                "session_id": session_id,
                "version": int(row["version"]),
                "saved_at": row["saved_at"],
                "state": json.loads(row["state_json"]),
            }
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to export state: {exc}") from exc

    def import_state(
        self,
        namespace: str,
        payload: dict[str, Any],
        *,
        session_id: str | None = None,
    ) -> AgentState:
        try:
            state_payload = payload.get("state", payload)
            state = _deserialize_state_data(state_payload)
            state.session_id = session_id if session_id is not None else state.session_id
            return self.save(namespace, state)
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to import state: {exc}") from exc

    def rollback(self, namespace: str, session_id: str | None, *, version: int) -> AgentState:
        exported = self.export_state(namespace, session_id, version=version)
        if exported is None:
            raise StateError(f"No saved state version {version} for session {session_id}")
        return self.import_state(namespace, exported["state"], session_id=session_id)

    def describe(self) -> dict[str, Any]:
        backend = "sqlite"
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                if self._uses_legacy_schema(connection):
                    backend = "sqlite_legacy"
        except Exception:
            pass
        return {
            "backend": backend,
            "path": self.path,
            "table_name": self.table_name,
            "schema_version": STATE_SCHEMA_VERSION,
            "shards": 1,
        }


class PostgreSQLStateStore(StateStore):
    """Persistent state store backed by PostgreSQL."""

    def __init__(self, dsn: str, table_name: str = "agent_state_versions") -> None:
        self.dsn = dsn
        self.table_name = safe_identifier(table_name)
        self._schema_ready = False

    def _connect(self) -> Any:
        try:
            return pooled_connect(self.dsn)
        except Exception as exc:
            raise StateBackendError(f"Failed to open PostgreSQL state backend: {exc}") from exc

    def _ensure_schema(self, connection: Any) -> None:
        if self._schema_ready:
            return
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    namespace TEXT NOT NULL,
                    session_id TEXT NOT NULL,
                    version INTEGER NOT NULL,
                    saved_at TIMESTAMPTZ NOT NULL,
                    state_json JSONB NOT NULL,
                    PRIMARY KEY (namespace, session_id, version)
                )
                """
            )
            cursor.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_latest
                ON {self.table_name} (namespace, session_id, version DESC)
                """
            )
        recorded_version = get_postgresql_schema_version(connection, STATE_SCHEMA_KIND)
        if recorded_version != STATE_SCHEMA_VERSION:
            set_postgresql_schema_version(connection, STATE_SCHEMA_KIND, STATE_SCHEMA_VERSION)
        connection.commit()
        self._schema_ready = True

    def _session_key(self, session_id: str | None) -> str:
        return session_id or ""

    def save(self, namespace: str, state: AgentState) -> AgentState:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                session_id = self._session_key(state.session_id)
                with connection.cursor() as cursor:
                    cursor.execute(
                        f"""
                        SELECT COALESCE(MAX(version), 0)
                        FROM {self.table_name}
                        WHERE namespace = %s AND session_id = %s
                        """,
                        (namespace, session_id),
                    )
                    current_version = int(cursor.fetchone()[0])
                    stored_state = _clone_state(state)
                    stored_state.version = current_version + 1
                    stored_state.updated_at = datetime.now(UTC)
                    cursor.execute(
                        f"""
                        INSERT INTO {self.table_name} (namespace, session_id, version, saved_at, state_json)
                        VALUES (%s, %s, %s, %s, %s::jsonb)
                        """,
                        (
                            namespace,
                            session_id,
                            stored_state.version,
                            stored_state.updated_at,
                            json.dumps(_serialize_state_data(stored_state), default=_json_default),
                        ),
                    )
                connection.commit()
                return _clone_state(stored_state)
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to persist state: {exc}") from exc

    def load(self, namespace: str, session_id: str | None) -> AgentState | None:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                with connection.cursor() as cursor:
                    cursor.execute(
                        f"""
                        SELECT version, saved_at, state_json
                        FROM {self.table_name}
                        WHERE namespace = %s AND session_id = %s
                        ORDER BY version DESC
                        LIMIT 1
                        """,
                        (namespace, self._session_key(session_id)),
                    )
                    row = cursor.fetchone()
            if row is None:
                return None
            state = _deserialize_state_data(row[2] if isinstance(row[2], dict) else json.loads(row[2]))
            state.version = int(row[0])
            state.updated_at = row[1] if isinstance(row[1], datetime) else datetime.fromisoformat(row[1])
            return state
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to load state: {exc}") from exc

    def history(self, namespace: str, session_id: str | None, *, limit: int = 50) -> list[dict[str, Any]]:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                with connection.cursor() as cursor:
                    cursor.execute(
                        f"""
                        SELECT version, saved_at
                        FROM {self.table_name}
                        WHERE namespace = %s AND session_id = %s
                        ORDER BY version DESC
                        LIMIT %s
                        """,
                        (namespace, self._session_key(session_id), limit),
                    )
                    rows = cursor.fetchall()
            return [
                {
                    "version": int(row[0]),
                    "saved_at": row[1].isoformat() if isinstance(row[1], datetime) else str(row[1]),
                    "session_id": session_id,
                }
                for row in rows
            ]
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to list state history: {exc}") from exc

    def export_state(
        self,
        namespace: str,
        session_id: str | None,
        *,
        version: int | None = None,
    ) -> dict[str, Any] | None:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                with connection.cursor() as cursor:
                    if version is None:
                        cursor.execute(
                            f"""
                            SELECT version, saved_at, state_json
                            FROM {self.table_name}
                            WHERE namespace = %s AND session_id = %s
                            ORDER BY version DESC
                            LIMIT 1
                            """,
                            (namespace, self._session_key(session_id)),
                        )
                    else:
                        cursor.execute(
                            f"""
                            SELECT version, saved_at, state_json
                            FROM {self.table_name}
                            WHERE namespace = %s AND session_id = %s AND version = %s
                            LIMIT 1
                            """,
                            (namespace, self._session_key(session_id), version),
                        )
                    row = cursor.fetchone()
            if row is None:
                return None
            state_json = row[2] if isinstance(row[2], dict) else json.loads(row[2])
            return {
                "agent_name": namespace,
                "session_id": session_id,
                "version": int(row[0]),
                "saved_at": row[1].isoformat() if isinstance(row[1], datetime) else str(row[1]),
                "state": state_json,
            }
        except StateBackendError:
            raise
        except Exception as exc:
            raise StateBackendError(f"Failed to export state: {exc}") from exc

    def import_state(
        self,
        namespace: str,
        payload: dict[str, Any],
        *,
        session_id: str | None = None,
    ) -> AgentState:
        state_payload = payload.get("state", payload)
        state = _deserialize_state_data(state_payload)
        state.session_id = session_id if session_id is not None else state.session_id
        return self.save(namespace, state)

    def rollback(self, namespace: str, session_id: str | None, *, version: int) -> AgentState:
        exported = self.export_state(namespace, session_id, version=version)
        if exported is None:
            raise StateError(f"No saved state version {version} for session {session_id}")
        return self.import_state(namespace, exported["state"], session_id=session_id)

    def describe(self) -> dict[str, Any]:
        return {
            "backend": "postgresql",
            "dsn": self.dsn,
            "table_name": self.table_name,
            "schema_version": STATE_SCHEMA_VERSION,
            "shards": 1,
        }


class ShardedStateStore(StateStore):
    """Routes state operations across multiple backend shards."""

    def __init__(self, shards: list[StateStore]) -> None:
        if not shards:
            raise ValueError("ShardedStateStore requires at least one shard")
        self._shards = shards

    def _shard_index(self, key: str) -> int:
        digest = hashlib.sha256(key.encode("utf-8")).hexdigest()
        return int(digest, 16) % len(self._shards)

    def _select_shard(self, session_id: str | None) -> StateStore:
        return self._shards[self._shard_index(session_id or "__default__")]

    def save(self, namespace: str, state: AgentState) -> AgentState:
        return self._select_shard(state.session_id).save(namespace, state)

    def load(self, namespace: str, session_id: str | None) -> AgentState | None:
        return self._select_shard(session_id).load(namespace, session_id)

    def history(self, namespace: str, session_id: str | None, *, limit: int = 50) -> list[dict[str, Any]]:
        return self._select_shard(session_id).history(namespace, session_id, limit=limit)

    def export_state(
        self,
        namespace: str,
        session_id: str | None,
        *,
        version: int | None = None,
    ) -> dict[str, Any] | None:
        return self._select_shard(session_id).export_state(namespace, session_id, version=version)

    def import_state(
        self,
        namespace: str,
        payload: dict[str, Any],
        *,
        session_id: str | None = None,
    ) -> AgentState:
        resolved_session_id = session_id
        if resolved_session_id is None:
            state_payload = payload.get("state", payload)
            resolved_session_id = state_payload.get("session_id")
        return self._select_shard(resolved_session_id).import_state(
            namespace,
            payload,
            session_id=session_id,
        )

    def rollback(self, namespace: str, session_id: str | None, *, version: int) -> AgentState:
        return self._select_shard(session_id).rollback(namespace, session_id, version=version)

    def describe(self) -> dict[str, Any]:
        return {
            "backend": "sharded",
            "shards": len(self._shards),
            "stores": [shard.describe() for shard in self._shards],
        }


def build_state_store(
    config: StateConfig | None,
    *,
    default_persistent: bool = False,
) -> StateStore:
    """Build a configured state store backend."""
    state_config = config or StateConfig()
    database_url = resolve_database_url(state_config.dsn)
    backend_name = state_config.backend

    if backend_name is None:
        if default_persistent:
            backend_name = "postgresql" if database_url else "sqlite"
        else:
            backend_name = "in_memory"

    if backend_name in {"in_memory", "memory", "inmemory"}:
        return InMemoryStateStore()
    if backend_name == "sqlite":
        if not state_config.path:
            raise ConfigError("SQLite state backend requires a non-empty path")
        return SQLiteStateStore(state_config.path, table_name=state_config.table_name)
    if backend_name in {"postgres", "postgresql"}:
        if not database_url:
            raise ConfigError(
                "PostgreSQL state backend requires a DSN or DATABASE_URL/LOGOS_ADK_DATABASE_URL"
            )
        return PostgreSQLStateStore(database_url, table_name=state_config.table_name)
    if backend_name == "sharded":
        if not state_config.shards:
            raise ConfigError("Sharded state backend requires a non-empty 'shards' list")
        return ShardedStateStore([build_state_store(shard, default_persistent=True) for shard in state_config.shards])

    raise ConfigError(f"Unsupported state backend: {backend_name}")


__all__ = [
    "StateStore",
    "InMemoryStateStore",
    "SQLiteStateStore",
    "PostgreSQLStateStore",
    "ShardedStateStore",
    "build_state_store",
]
