"""
停等账仓储壳：suspend_wait（waitVersion）。

本文件职责：登记/关闭停点；与 Graph interrupt、H5 token 对齐的版本号。
禁止：用本表冒充主账六态；禁止用户文字伪造 Worker 完成。
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Optional

from agent.runtime_db import assert_safe_database, connect_psycopg

logger = logging.getLogger(__name__)


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


@dataclass
class SuspendRecord:
    """一条停等记录。"""

    thread_id: str
    wait_version: int
    pause_name: str = ""
    inquiry_no: str = ""
    status: str = "open"
    payload: dict[str, Any] = field(default_factory=dict)


class MemorySuspendStore:
    """内存停等壳。"""

    def __init__(self) -> None:
        self._rows: dict[tuple[str, int], SuspendRecord] = {}
        self._lock = threading.Lock()

    def open_wait(self, record: SuspendRecord) -> None:
        key = (record.thread_id, record.wait_version)
        with self._lock:
            self._rows[key] = SuspendRecord(
                thread_id=record.thread_id,
                wait_version=record.wait_version,
                pause_name=record.pause_name,
                inquiry_no=record.inquiry_no,
                status="open",
                payload=dict(record.payload),
            )
        logger.info(
            "suspend.memory.open thread=%s wait=%s pause=%s",
            record.thread_id,
            record.wait_version,
            record.pause_name,
        )

    def get(self, thread_id: str, wait_version: int) -> Optional[SuspendRecord]:
        with self._lock:
            rec = self._rows.get((thread_id, wait_version))
            if rec is None:
                return None
            return SuspendRecord(
                thread_id=rec.thread_id,
                wait_version=rec.wait_version,
                pause_name=rec.pause_name,
                inquiry_no=rec.inquiry_no,
                status=rec.status,
                payload=dict(rec.payload),
            )

    def close_wait(self, thread_id: str, wait_version: int) -> bool:
        with self._lock:
            rec = self._rows.get((thread_id, wait_version))
            if rec is None:
                return False
            rec.status = "closed"
            return True


class PostgresSuspendStore:
    """PG 停等壳（表 suspend_wait）。"""

    def __init__(self, dsn: str) -> None:
        assert_safe_database(dsn)
        self._dsn = dsn

    def open_wait(self, record: SuspendRecord) -> None:
        import json

        with connect_psycopg(self._dsn) as conn:
            conn.execute(
                """
                INSERT INTO suspend_wait (
                    thread_id, wait_version, pause_name, inquiry_no, status, payload
                ) VALUES (%s, %s, %s, %s, 'open', %s::jsonb)
                ON CONFLICT (thread_id, wait_version) DO UPDATE SET
                    pause_name = EXCLUDED.pause_name,
                    inquiry_no = EXCLUDED.inquiry_no,
                    status = 'open',
                    payload = EXCLUDED.payload,
                    closed_at = NULL
                """,
                (
                    record.thread_id,
                    record.wait_version,
                    record.pause_name,
                    record.inquiry_no,
                    json.dumps(record.payload, ensure_ascii=False),
                ),
            )
            conn.commit()
        logger.info(
            "suspend.pg.open thread=%s wait=%s",
            record.thread_id,
            record.wait_version,
        )

    def get(self, thread_id: str, wait_version: int) -> Optional[SuspendRecord]:
        with connect_psycopg(self._dsn) as conn:
            row = conn.execute(
                """
                SELECT thread_id, wait_version, pause_name, inquiry_no, status, payload
                FROM suspend_wait
                WHERE thread_id = %s AND wait_version = %s
                """,
                (thread_id, wait_version),
            ).fetchone()
        if not row:
            return None
        return SuspendRecord(
            thread_id=str(row[0]),
            wait_version=int(row[1]),
            pause_name=str(row[2] or ""),
            inquiry_no=str(row[3] or ""),
            status=str(row[4] or ""),
            payload=dict(row[5] or {}),
        )

    def close_wait(self, thread_id: str, wait_version: int) -> bool:
        with connect_psycopg(self._dsn) as conn:
            cur = conn.execute(
                """
                UPDATE suspend_wait
                SET status = 'closed', closed_at = %s
                WHERE thread_id = %s AND wait_version = %s AND status = 'open'
                """,
                (_utcnow(), thread_id, wait_version),
            )
            conn.commit()
            return cur.rowcount > 0


def create_suspend_store(*, dsn: str = "", backend: str = "auto"):
    """工厂：Memory 或 PG。"""
    if backend == "memory" or not dsn:
        return MemorySuspendStore()
    return PostgresSuspendStore(dsn)