"""DuckDB-backed runtime state store for pipeline execution."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import json
from pathlib import Path
from typing import Any, Optional
import uuid

import duckdb


_LOCAL_TIMEZONE = datetime.now().astimezone().tzinfo or timezone.utc

@dataclass(frozen=True)
class PendingRunEvent:
    event_id: str
    pipeline_id: str
    event_type: str
    source: str
    payload: Any
    created_at: Optional[datetime]
    dedupe_key: Optional[str]
    source_event_id: Optional[str]
    defer_count: int
    next_attempt_at: Optional[datetime]


def _normalize_timestamp(value: Optional[datetime]) -> Optional[datetime]:
    if value is None:
        return None
    if value.tzinfo is None:
        return value.replace(tzinfo=_LOCAL_TIMEZONE).astimezone(timezone.utc)
    return value.astimezone(timezone.utc)


def _encode_payload(payload: Any) -> str:
    if payload is None:
        payload = {}
    return json.dumps(payload, sort_keys=True, separators=(",", ":"))


def _decode_payload(payload: Any) -> Any:
    if payload is None:
        return {}
    if isinstance(payload, str):
        try:
            return json.loads(payload)
        except Exception:
            return payload
    return payload


def _derive_session_id(now: datetime) -> str:
    now_utc = _normalize_timestamp(now) or datetime.now(timezone.utc)
    return f"session-{now_utc.strftime('%Y%m%d')}"


def _next_session_run_index(
    con: duckdb.DuckDBPyConnection,
    *,
    pipeline_id: str,
    session_id: str,
) -> int:
    row = con.execute(
        """
        SELECT COALESCE(MAX(session_run_index), -1)
        FROM pipeline_runs
        WHERE pipeline_id = ? AND session_id = ?
        """,
        [pipeline_id, session_id],
    ).fetchone()
    if not row or row[0] is None:
        return 0
    return int(row[0]) + 1


class PipelineStateStore:
    """
    Pipeline runtime state store (non-RDF).

    This store manages the high-churn pipeline execution state in
    `cf-pipeline-states.duckdb`.
    """

    def __init__(self, db_path: Path, semantics_dir: Path, read_only: bool = False):
        self.db_path = Path(db_path)
        self.semantics_dir = Path(semantics_dir)
        self.read_only = bool(read_only)

    def _connect(self, *, read_only: Optional[bool] = None) -> duckdb.DuckDBPyConnection:
        ro = self.read_only if read_only is None else bool(read_only)
        return duckdb.connect(str(self.db_path), read_only=ro)

    def init_schema(self) -> None:
        """Create runtime-state tables and indexes when missing."""
        if not self.read_only:
            self.semantics_dir.mkdir(parents=True, exist_ok=True)

        def _table_columns(con: duckdb.DuckDBPyConnection, table_name: str) -> set[str]:
            rows = con.execute(f"PRAGMA table_info('{table_name}')").fetchall()
            cols: set[str] = set()
            for row in rows:
                # DuckDB PRAGMA table_info columns:
                # cid, name, type, notnull, dflt_value, pk
                if len(row) >= 2 and isinstance(row[1], str):
                    cols.add(row[1])
            return cols

        con = self._connect(read_only=False)
        try:
            con.execute(
                """
                CREATE TABLE IF NOT EXISTS pipeline_state (
                  pipeline_id TEXT PRIMARY KEY,
                  current_rev INTEGER NOT NULL,
                  desired_state TEXT NOT NULL,
                  concurrency_limit INTEGER DEFAULT 1,
                  priority INTEGER DEFAULT 0,
                  last_run_id TEXT,
                  last_run_at TIMESTAMP,
                  next_run_at TIMESTAMP,
                  created_at TIMESTAMP DEFAULT current_timestamp,
                  updated_at TIMESTAMP DEFAULT current_timestamp
                )
                """
            )
            con.execute(
                """
                CREATE TABLE IF NOT EXISTS pipeline_runs (
                  run_id TEXT PRIMARY KEY,
                  pipeline_id TEXT NOT NULL,
                  session_id TEXT NOT NULL,
                  session_run_index INTEGER NOT NULL,
                  pipeline_rev INTEGER NOT NULL,
                  status TEXT NOT NULL,
                  triggered_by_event_id TEXT,
                  created_at TIMESTAMP DEFAULT current_timestamp,
                  started_at TIMESTAMP,
                  finished_at TIMESTAMP,
                  termination_reason TEXT,
                  worker_id TEXT,
                  attempt INTEGER DEFAULT 0
                )
                """
            )
            con.execute(
                """
                CREATE TABLE IF NOT EXISTS run_events (
                  event_id TEXT PRIMARY KEY,
                  pipeline_id TEXT NOT NULL,
                  run_id TEXT,
                  event_type TEXT NOT NULL,
                  source TEXT NOT NULL DEFAULT 'system',
                  source_event_id TEXT,
                  payload JSON,
                  status TEXT NOT NULL DEFAULT 'pending',
                  created_at TIMESTAMP DEFAULT current_timestamp,
                  updated_at TIMESTAMP DEFAULT current_timestamp,
                  next_attempt_at TIMESTAMP,
                  defer_count INTEGER DEFAULT 0,
                  defer_reason TEXT,
                  last_error TEXT,
                  consumed_at TIMESTAMP,
                  dedupe_key TEXT
                )
                """
            )
            con.execute(
                """
                CREATE TABLE IF NOT EXISTS run_queue (
                  queue_id TEXT PRIMARY KEY,
                  pipeline_id TEXT NOT NULL,
                  run_id TEXT NOT NULL,
                  available_at TIMESTAMP NOT NULL,
                  priority INTEGER DEFAULT 0,
                  locked_by TEXT,
                  lock_expires_at TIMESTAMP,
                  attempts INTEGER DEFAULT 0,
                  last_error TEXT,
                  created_at TIMESTAMP DEFAULT current_timestamp,
                  updated_at TIMESTAMP DEFAULT current_timestamp
                )
                """
            )

            con.execute(
                """
                CREATE UNIQUE INDEX IF NOT EXISTS idx_run_events_dedupe
                ON run_events(dedupe_key)
                """
            )
            con.execute(
                """
                CREATE UNIQUE INDEX IF NOT EXISTS idx_run_queue_run
                ON run_queue(run_id)
                """
            )
            con.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_pipeline_runs_pipeline
                ON pipeline_runs(pipeline_id, status)
                """
            )
            con.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_run_events_pipeline
                ON run_events(pipeline_id)
                """
            )
            con.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_run_events_run
                ON run_events(run_id)
                """
            )
            con.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_run_events_pending
                ON run_events(pipeline_id, status, next_attempt_at, created_at)
                """
            )
            con.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_run_queue_available
                ON run_queue(available_at, priority)
                """
            )

            pipeline_runs_cols = _table_columns(con, "pipeline_runs")
            if "session_id" not in pipeline_runs_cols:
                con.execute("ALTER TABLE pipeline_runs ADD COLUMN session_id TEXT")
            if "session_run_index" not in pipeline_runs_cols:
                con.execute("ALTER TABLE pipeline_runs ADD COLUMN session_run_index INTEGER")

            run_events_cols = _table_columns(con, "run_events")
            if "source" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN source TEXT DEFAULT 'system'")
            if "source_event_id" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN source_event_id TEXT")
            if "status" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN status TEXT DEFAULT 'pending'")
            if "updated_at" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN updated_at TIMESTAMP DEFAULT current_timestamp")
            if "next_attempt_at" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN next_attempt_at TIMESTAMP")
            if "defer_count" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN defer_count INTEGER DEFAULT 0")
            if "defer_reason" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN defer_reason TEXT")
            if "last_error" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN last_error TEXT")
            if "consumed_at" not in run_events_cols:
                con.execute("ALTER TABLE run_events ADD COLUMN consumed_at TIMESTAMP")

            con.execute(
                """
                UPDATE pipeline_runs
                SET session_id = COALESCE(session_id, pipeline_id || ':legacy-session')
                WHERE session_id IS NULL OR TRIM(session_id) = ''
                """
            )
            con.execute(
                """
                UPDATE pipeline_runs AS pr
                SET session_run_index = ranked.session_run_index
                FROM (
                  SELECT
                    run_id,
                    ROW_NUMBER() OVER (
                      PARTITION BY pipeline_id, session_id
                      ORDER BY created_at, run_id
                    ) - 1 AS session_run_index
                  FROM pipeline_runs
                ) AS ranked
                WHERE pr.run_id = ranked.run_id
                  AND pr.session_run_index IS NULL
                """
            )
            con.execute(
                """
                UPDATE run_events
                SET source = COALESCE(NULLIF(TRIM(source), ''), 'system')
                WHERE source IS NULL OR TRIM(source) = ''
                """
            )
            con.execute(
                """
                UPDATE run_events
                SET status = CASE
                  WHEN run_id IS NULL THEN 'pending'
                  ELSE 'consumed'
                END
                WHERE status IS NULL OR TRIM(status) = ''
                """
            )
            con.execute(
                """
                UPDATE run_events
                SET updated_at = COALESCE(updated_at, created_at, current_timestamp)
                WHERE updated_at IS NULL
                """
            )
            con.execute(
                """
                UPDATE run_events
                SET consumed_at = COALESCE(consumed_at, updated_at, created_at, current_timestamp)
                WHERE run_id IS NOT NULL AND consumed_at IS NULL
                """
            )
            con.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_pipeline_runs_session
                ON pipeline_runs(pipeline_id, session_id, session_run_index)
                """
            )
        finally:
            con.close()

    def insert_event(
        self,
        *,
        pipeline_id: str,
        event_type: str,
        payload: Any = None,
        source: str = "external",
        dedupe_key: Optional[str] = None,
        source_event_id: Optional[str] = None,
        con: Optional[duckdb.DuckDBPyConnection] = None,
    ) -> tuple[str, bool]:
        own_connection = con is None
        db = con or self._connect(read_only=False)
        norm_source = (source or "external").strip() or "external"
        norm_dedupe = dedupe_key.strip() if isinstance(dedupe_key, str) and dedupe_key.strip() else None
        event_id = str(uuid.uuid4())
        payload_json = _encode_payload(payload)
        try:
            if norm_dedupe:
                existing = db.execute(
                    "SELECT event_id FROM run_events WHERE dedupe_key = ? LIMIT 1",
                    [norm_dedupe],
                ).fetchone()
                if existing and existing[0]:
                    return str(existing[0]), False

            try:
                db.execute(
                    """
                    INSERT INTO run_events (
                      event_id,
                      pipeline_id,
                      run_id,
                      event_type,
                      source,
                      source_event_id,
                      payload,
                      status,
                      created_at,
                      updated_at,
                      dedupe_key
                    )
                    VALUES (
                      ?,
                      ?,
                      NULL,
                      ?,
                      ?,
                      ?,
                      CAST(? AS JSON),
                      'pending',
                      current_timestamp,
                      current_timestamp,
                      ?
                    )
                    """,
                    [
                        event_id,
                        pipeline_id,
                        event_type,
                        norm_source,
                        source_event_id,
                        payload_json,
                        norm_dedupe,
                    ],
                )
            except duckdb.ConstraintException:
                if norm_dedupe:
                    existing = db.execute(
                        "SELECT event_id FROM run_events WHERE dedupe_key = ? LIMIT 1",
                        [norm_dedupe],
                    ).fetchone()
                    if existing and existing[0]:
                        return str(existing[0]), False
                raise
            return event_id, True
        finally:
            if own_connection:
                db.close()

    def fetch_next_pending_event(
        self,
        *,
        pipeline_id: Optional[str] = None,
        now: Optional[datetime] = None,
        con: Optional[duckdb.DuckDBPyConnection] = None,
    ) -> Optional[PendingRunEvent]:
        own_connection = con is None
        db = con or self._connect(read_only=False)
        now_utc = _normalize_timestamp(now) or datetime.now(timezone.utc)
        try:
            query = """
                SELECT
                  event_id,
                  pipeline_id,
                  event_type,
                  source,
                  source_event_id,
                  payload,
                  created_at,
                  dedupe_key,
                  defer_count,
                  next_attempt_at
                FROM run_events
                WHERE run_id IS NULL
                  AND status = 'pending'
                  AND (next_attempt_at IS NULL OR next_attempt_at <= ?)
            """
            params: list[Any] = [now_utc]
            if pipeline_id:
                query += " AND pipeline_id = ?"
                params.append(pipeline_id)
            query += """
                ORDER BY COALESCE(next_attempt_at, created_at), created_at, event_id
                LIMIT 1
            """
            row = db.execute(query, params).fetchone()
            if not row:
                return None
            return PendingRunEvent(
                event_id=str(row[0]),
                pipeline_id=str(row[1]),
                event_type=str(row[2]),
                source=str(row[3] or "system"),
                source_event_id=str(row[4]) if row[4] is not None else None,
                payload=_decode_payload(row[5]),
                created_at=_normalize_timestamp(row[6]),
                dedupe_key=str(row[7]) if row[7] is not None else None,
                defer_count=int(row[8] or 0),
                next_attempt_at=_normalize_timestamp(row[9]),
            )
        finally:
            if own_connection:
                db.close()

    def defer_event(
        self,
        *,
        event_id: str,
        reason: str,
        next_attempt_at: datetime,
        last_error: Optional[str] = None,
        con: Optional[duckdb.DuckDBPyConnection] = None,
    ) -> bool:
        own_connection = con is None
        db = con or self._connect(read_only=False)
        next_attempt = _normalize_timestamp(next_attempt_at) or datetime.now(timezone.utc)
        try:
            row = db.execute(
                """
                SELECT
                  pipeline_id,
                  event_type,
                  source,
                  source_event_id,
                  payload,
                  created_at,
                  dedupe_key,
                  COALESCE(defer_count, 0)
                FROM run_events
                WHERE event_id = ?
                  AND run_id IS NULL
                LIMIT 1
                """,
                [event_id],
            ).fetchone()
            if not row:
                return False
            db.execute(
                "DELETE FROM run_events WHERE event_id = ? AND run_id IS NULL",
                [event_id],
            )
            db.execute(
                """
                INSERT INTO run_events (
                  event_id,
                  pipeline_id,
                  run_id,
                  event_type,
                  source,
                  source_event_id,
                  payload,
                  status,
                  created_at,
                  updated_at,
                  next_attempt_at,
                  defer_count,
                  defer_reason,
                  last_error,
                  consumed_at,
                  dedupe_key
                )
                VALUES (
                  ?,
                  ?,
                  NULL,
                  ?,
                  ?,
                  ?,
                  CAST(? AS JSON),
                  'pending',
                  ?,
                  current_timestamp,
                  ?,
                  ?,
                  ?,
                  ?,
                  NULL,
                  ?
                )
                """,
                [
                    event_id,
                    str(row[0]),
                    str(row[1]),
                    str(row[2] or "system"),
                    str(row[3]) if row[3] is not None else None,
                    _encode_payload(_decode_payload(row[4])),
                    row[5],
                    next_attempt,
                    int(row[7] or 0) + 1,
                    reason,
                    last_error,
                    str(row[6]) if row[6] is not None else None,
                ],
            )
            return True
        finally:
            if own_connection:
                db.close()

    def consume_event_to_run(
        self,
        *,
        event_id: str,
        pipeline_rev: int,
        priority: int,
        now: Optional[datetime] = None,
        con: Optional[duckdb.DuckDBPyConnection] = None,
    ) -> Optional[tuple[str, str, str]]:
        own_connection = con is None
        db = con or self._connect(read_only=False)
        now_utc = _normalize_timestamp(now) or datetime.now(timezone.utc)
        run_id = str(uuid.uuid4())
        queue_id = str(uuid.uuid4())
        try:
            candidate = db.execute(
                """
                SELECT
                  pipeline_id,
                  event_type,
                  source,
                  source_event_id,
                  payload,
                  created_at,
                  dedupe_key
                FROM run_events
                WHERE event_id = ?
                  AND run_id IS NULL
                  AND status = 'pending'
                LIMIT 1
                """,
                [event_id],
            ).fetchone()
            if not candidate or candidate[0] is None:
                return None
            pipeline_id = str(candidate[0])

            db.execute(
                "DELETE FROM run_events WHERE event_id = ? AND run_id IS NULL AND status = 'pending'",
                [event_id],
            )
            still_pending = db.execute(
                """
                SELECT 1
                FROM run_events
                WHERE event_id = ?
                  AND run_id IS NULL
                  AND status = 'pending'
                LIMIT 1
                """,
                [event_id],
            ).fetchone()
            if still_pending:
                return None

            db.execute(
                """
                INSERT INTO run_events (
                  event_id,
                  pipeline_id,
                  run_id,
                  event_type,
                  source,
                  source_event_id,
                  payload,
                  status,
                  created_at,
                  updated_at,
                  next_attempt_at,
                  defer_count,
                  defer_reason,
                  last_error,
                  consumed_at,
                  dedupe_key
                )
                VALUES (
                  ?,
                  ?,
                  ?,
                  ?,
                  ?,
                  ?,
                  CAST(? AS JSON),
                  'consumed',
                  ?,
                  current_timestamp,
                  NULL,
                  0,
                  NULL,
                  NULL,
                  current_timestamp,
                  ?
                )
                """,
                [
                    event_id,
                    pipeline_id,
                    run_id,
                    str(candidate[1]),
                    str(candidate[2] or "system"),
                    str(candidate[3]) if candidate[3] is not None else None,
                    _encode_payload(_decode_payload(candidate[4])),
                    candidate[5],
                    str(candidate[6]) if candidate[6] is not None else None,
                ],
            )
            consumed_check = db.execute(
                """
                SELECT pipeline_id
                FROM run_events
                WHERE event_id = ? AND run_id = ?
                LIMIT 1
                """,
                [event_id, run_id],
            ).fetchone()
            if not consumed_check:
                return None

            session_id = _derive_session_id(now_utc)
            session_run_index = _next_session_run_index(
                db,
                pipeline_id=pipeline_id,
                session_id=session_id,
            )

            db.execute(
                """
                INSERT INTO pipeline_runs (
                  run_id,
                  pipeline_id,
                  session_id,
                  session_run_index,
                  pipeline_rev,
                  status,
                  triggered_by_event_id,
                  created_at
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, current_timestamp)
                """,
                [
                    run_id,
                    pipeline_id,
                    session_id,
                    int(session_run_index),
                    int(pipeline_rev),
                    "queued",
                    event_id,
                ],
            )
            db.execute(
                """
                INSERT INTO run_queue (
                  queue_id,
                  pipeline_id,
                  run_id,
                  available_at,
                  priority,
                  created_at,
                  updated_at
                )
                VALUES (?, ?, ?, ?, ?, current_timestamp, current_timestamp)
                """,
                [queue_id, pipeline_id, run_id, now_utc, int(priority)],
            )
            registration = db.execute(
                """
                SELECT
                  re.run_id,
                  pr.triggered_by_event_id,
                  rq.queue_id
                FROM run_events AS re
                JOIN pipeline_runs AS pr
                  ON pr.run_id = re.run_id
                JOIN run_queue AS rq
                  ON rq.run_id = pr.run_id
                WHERE re.event_id = ?
                  AND re.run_id = ?
                LIMIT 1
                """,
                [event_id, run_id],
            ).fetchone()
            if not registration:
                return None
            if str(registration[1] or "") != event_id:
                return None
            if str(registration[2] or "") != queue_id:
                return None
            return run_id, queue_id, pipeline_id
        finally:
            if own_connection:
                db.close()
