"""Cross-run pipeline state persistence.

Provides a :class:`StateStore` protocol and two concrete implementations
(:class:`FileStateStore`, :class:`SqliteStateStore`) for persisting watermarks
and run metadata between pipeline invocations.
"""

from __future__ import annotations

import json
import logging
import sqlite3
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Protocol, runtime_checkable

logger = logging.getLogger(__name__)


@dataclass(frozen=True, slots=True)
class PipelineState:
    """State persisted between pipeline runs."""

    pipeline_name: str
    watermark: str | None = None
    last_run_id: str | None = None
    last_run_at: str | None = None  # ISO 8601
    run_count: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Serialize state to a plain dictionary."""
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PipelineState:
        """Deserialize state from a plain dictionary."""
        return cls(
            pipeline_name=data["pipeline_name"],
            watermark=data.get("watermark"),
            last_run_id=data.get("last_run_id"),
            last_run_at=data.get("last_run_at"),
            run_count=data.get("run_count", 0),
            metadata=data.get("metadata") or {},
        )


@runtime_checkable
class StateStore(Protocol):
    """Protocol for pipeline state persistence."""

    def get(self, pipeline_name: str) -> PipelineState | None:
        """Retrieve state for a pipeline, or *None* if not found."""
        ...

    def save(self, state: PipelineState) -> None:
        """Persist pipeline state."""
        ...

    def delete(self, pipeline_name: str) -> None:
        """Delete state for a pipeline."""
        ...

    def list_pipelines(self) -> list[str]:
        """List pipeline names that have persisted state."""
        ...


class FileStateStore:
    """Store pipeline state as JSON files in a directory.

    One file per pipeline: ``{pipeline_name}.json``.
    """

    def __init__(self, path: str | Path = "./.state") -> None:
        self._path = Path(path)

    def get(self, pipeline_name: str) -> PipelineState | None:
        """Retrieve state for *pipeline_name*."""
        state_file = self._state_file(pipeline_name)
        if not state_file.exists():
            return None
        data = json.loads(state_file.read_text(encoding="utf-8"))
        return PipelineState.from_dict(data)

    def save(self, state: PipelineState) -> None:
        """Persist *state* to a JSON file."""
        self._path.mkdir(parents=True, exist_ok=True)
        state_file = self._state_file(state.pipeline_name)
        state_file.write_text(
            json.dumps(state.to_dict(), indent=2, default=str),
            encoding="utf-8",
        )
        logger.debug("Saved state for %s", state.pipeline_name)

    def delete(self, pipeline_name: str) -> None:
        """Delete state file for *pipeline_name*."""
        state_file = self._state_file(pipeline_name)
        if state_file.exists():
            state_file.unlink()
            logger.debug("Deleted state for %s", pipeline_name)

    def list_pipelines(self) -> list[str]:
        """List pipeline names with persisted state."""
        if not self._path.exists():
            return []
        return sorted(p.stem for p in self._path.glob("*.json"))

    def _state_file(self, pipeline_name: str) -> Path:
        """Return the JSON file path for *pipeline_name*."""
        return self._path / f"{pipeline_name}.json"


class SqliteStateStore:
    """Store pipeline state in a SQLite database table."""

    _TABLE = "pipeline_state"

    def __init__(self, path: str | Path = "./.state/state.db") -> None:
        self._path = Path(path)
        self._ensure_table()

    def _connect(self) -> sqlite3.Connection:
        """Open a connection, creating parent dirs as needed."""
        self._path.parent.mkdir(parents=True, exist_ok=True)
        return sqlite3.connect(str(self._path))

    def _ensure_table(self) -> None:
        """Create the state table if it does not exist."""
        conn = self._connect()
        try:
            conn.execute(
                f"CREATE TABLE IF NOT EXISTS {self._TABLE} ("
                "  pipeline_name TEXT PRIMARY KEY,"
                "  watermark TEXT,"
                "  last_run_id TEXT,"
                "  last_run_at TEXT,"
                "  run_count INTEGER DEFAULT 0,"
                "  metadata TEXT DEFAULT '{}'"
                ")"
            )
            conn.commit()
        finally:
            conn.close()

    def get(self, pipeline_name: str) -> PipelineState | None:
        """Retrieve state for *pipeline_name*."""
        conn = self._connect()
        try:
            row = conn.execute(
                f"SELECT pipeline_name, watermark, last_run_id, last_run_at, "
                f"run_count, metadata FROM {self._TABLE} WHERE pipeline_name = ?",
                (pipeline_name,),
            ).fetchone()
            if row is None:
                return None
            return PipelineState(
                pipeline_name=row[0],
                watermark=row[1],
                last_run_id=row[2],
                last_run_at=row[3],
                run_count=row[4],
                metadata=json.loads(row[5]) if row[5] else {},
            )
        finally:
            conn.close()

    def save(self, state: PipelineState) -> None:
        """Persist *state* to the SQLite database."""
        conn = self._connect()
        try:
            conn.execute(
                f"INSERT OR REPLACE INTO {self._TABLE} "
                "(pipeline_name, watermark, last_run_id, last_run_at, run_count, metadata) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                (
                    state.pipeline_name,
                    state.watermark,
                    state.last_run_id,
                    state.last_run_at,
                    state.run_count,
                    json.dumps(state.metadata, default=str),
                ),
            )
            conn.commit()
            logger.debug("Saved state for %s", state.pipeline_name)
        finally:
            conn.close()

    def delete(self, pipeline_name: str) -> None:
        """Delete state for *pipeline_name*."""
        conn = self._connect()
        try:
            conn.execute(
                f"DELETE FROM {self._TABLE} WHERE pipeline_name = ?",
                (pipeline_name,),
            )
            conn.commit()
            logger.debug("Deleted state for %s", pipeline_name)
        finally:
            conn.close()

    def list_pipelines(self) -> list[str]:
        """List pipeline names with persisted state."""
        conn = self._connect()
        try:
            rows = conn.execute(
                f"SELECT pipeline_name FROM {self._TABLE} ORDER BY pipeline_name"
            ).fetchall()
            return [r[0] for r in rows]
        finally:
            conn.close()


def create_state_store(config: dict[str, Any]) -> StateStore:
    """Create a state store from a configuration dict.

    Args:
        config: Dict with optional ``backend`` (``file`` or ``sqlite``) and
            ``path`` keys.

    Returns:
        A :class:`StateStore` implementation.

    Raises:
        ValueError: If the backend is unknown.
    """
    backend = config.get("backend", "file")
    if backend == "file":
        return FileStateStore(path=config.get("path", "./.state"))
    elif backend == "sqlite":
        return SqliteStateStore(path=config.get("path", "./.state/state.db"))
    else:
        raise ValueError(f"Unknown state backend: {backend}")
