YG_FT/backend/app/modules/data_process/store.py

from __future__ import annotations

import hashlib
import json
import uuid
from contextlib import contextmanager
from datetime import date, datetime, timezone
from functools import lru_cache
from pathlib import Path
from typing import Any, Iterator, Sequence

import psycopg
from psycopg.rows import dict_row

from app.core.config import get_settings
from app.modules.data_process.algorithms import estimate_token_count, stable_split

TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"}
EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"}


class DataProcessStoreError(RuntimeError):
    pass


class NotFoundError(DataProcessStoreError):
    pass


class ConflictError(DataProcessStoreError):
    pass


class InvalidStateError(DataProcessStoreError):
    pass


def utcnow() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def new_id(prefix: str) -> str:
    return f"{prefix}_{uuid.uuid4().hex[:20]}"


def json_dumps(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def _database_url(value: str) -> str:
    return value.replace("postgresql+psycopg://", "postgresql://")


def _json_value(value: Any, default: Any) -> Any:
    if value is None or value == "":
        return default
    if isinstance(value, (dict, list)):
        return value
    try:
        return json.loads(value)
    except (TypeError, json.JSONDecodeError):
        return default


def _serialize_value(value: Any) -> Any:
    if isinstance(value, (datetime, date)):
        return value.isoformat().replace("+00:00", "Z")
    return value


def _source_storage_descriptor(
    payload: dict[str, Any],
    task_id: str,
    file_id: str,
) -> tuple[str, dict[str, Any]]:
    storage_object_id = str(
        payload.get("storage_object_id")
        or f"db://data-process/{task_id}/{file_id}/v1"
    )
    expected_local_prefix = f"local://data-process/{task_id}/{file_id}/v1/"
    expected_database_reference = f"db://data-process/{task_id}/{file_id}/v1"
    if storage_object_id.startswith(expected_local_prefix) and len(storage_object_id) > len(
        expected_local_prefix
    ):
        storage_backend = "local"
    elif storage_object_id == expected_database_reference:
        storage_backend = "database"
    elif storage_object_id.startswith(("local://data-process/", "db://data-process/")):
        raise DataProcessStoreError("source storage object owner mismatch")
    else:
        raise DataProcessStoreError("unsupported source storage object reference")
    metadata = {
        **(payload.get("metadata") or {}),
        "storage_backend": storage_backend,
    }
    return storage_object_id, metadata


def _decode_row(row: dict[str, Any] | None) -> dict[str, Any] | None:
    if row is None:
        return None
    item = {key: _serialize_value(value) for key, value in row.items()}
    for key, default in {
        "config": {},
        "metadata": {},
        "quality_score": {},
        "versions": [],
    }.items():
        if key in item:
            item[key] = _json_value(item[key], default)
    return item


class DataProcessStore:
    """数据处理持久层。

    构造函数不会连接数据库或执行迁移。部署方必须显式执行 002 SQL，
    或在受控的管理命令中调用 :meth:`ensure_schema`，避免应用启动时
    修改远程数据库。
    """

    def __init__(self, database_url: str | None = None) -> None:
        self.database_url = _database_url(database_url or get_settings().database_url)

    @contextmanager
    def connect(self) -> Iterator[psycopg.Connection[dict[str, Any]]]:
        with psycopg.connect(self.database_url, row_factory=dict_row) as conn:
            try:
                yield conn
                conn.commit()
            except Exception:
                conn.rollback()
                raise

    def ensure_schema(self) -> None:
        """显式安装数据处理表；API 路由和应用启动流程不会调用此方法。"""
        schema_path = Path(__file__).resolve().parents[2] / "db" / "sql" / "002_data_process.sql"
        sql = schema_path.read_text(encoding="utf-8")
        with self.connect() as conn:
            with conn.cursor() as cursor:
                cursor.execute(sql)

    def list_tasks(
        self,
        *,
        page: int = 1,
        page_size: int = 20,
        keyword: str | None = None,
        status: str | None = None,
        process_type: str | None = None,
        tenant_id: str | None = None,
        project_id: str | None = None,
    ) -> dict[str, Any]:
        clauses = ["deleted_at IS NULL"]
        params: list[Any] = []
        if keyword:
            clauses.append("(name ILIKE %s OR COALESCE(description, '') ILIKE %s)")
            pattern = f"%{keyword.strip()}%"
            params.extend([pattern, pattern])
        if status:
            clauses.append("status = %s")
            params.append(status)
        if process_type:
            clauses.append("process_type = %s")
            params.append(process_type)
        if tenant_id:
            clauses.append("tenant_id = %s")
            params.append(tenant_id)
        if project_id:
            clauses.append("project_id = %s")
            params.append(project_id)
        where = " AND ".join(clauses)
        with self.connect() as conn:
            total = conn.execute(
                f"SELECT COUNT(*) AS count FROM data_process_tasks WHERE {where}", params
            ).fetchone()["count"]
            rows = conn.execute(
                f"""
                SELECT * FROM data_process_tasks
                WHERE {where}
                ORDER BY created_at DESC, id DESC
                LIMIT %s OFFSET %s
                """,
                [*params, page_size, (page - 1) * page_size],
            ).fetchall()
        return {
            "items": [_decode_row(row) for row in rows],
            "total": int(total),
            "page": page,
            "page_size": page_size,
        }

    def create_task(self, payload: dict[str, Any]) -> dict[str, Any]:
        task_id = new_id("dpt")
        now = utcnow()
        try:
            with self.connect() as conn:
                row = conn.execute(
                    """
                    INSERT INTO data_process_tasks
                    (id, name, description, status, process_type, source_dataset_id, config,
                     progress, tenant_id, project_id, owner_id, created_by, updated_by,
                     created_at, updated_at)
                    VALUES (%s, %s, %s, 'pending', %s, %s, %s, 0, %s, %s, %s, %s, %s, %s, %s)
                    RETURNING *
                    """,
                    (
                        task_id,
                        payload["name"],
                        payload.get("description") or "",
                        payload["process_type"],
                        payload.get("source_dataset_id"),
                        json_dumps(payload.get("config") or {}),
                        payload.get("tenant_id"),
                        payload.get("project_id"),
                        payload.get("owner_id"),
                        payload.get("created_by"),
                        payload.get("created_by"),
                        now,
                        now,
                    ),
                ).fetchone()
        except psycopg.errors.UniqueViolation as exc:
            raise ConflictError("data process task name already exists") from exc
        return _decode_row(row) or {}

    def get_task(self, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
        lock = " FOR UPDATE" if for_update else ""
        with self.connect() as conn:
            row = conn.execute(
                f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
                (task_id,),
            ).fetchone()
        if not row:
            raise NotFoundError("data process task not found")
        return _decode_row(row) or {}

    def _task_in_connection(
        self,
        conn: psycopg.Connection[dict[str, Any]],
        task_id: str,
        *,
        for_update: bool = False,
    ) -> dict[str, Any]:
        lock = " FOR UPDATE" if for_update else ""
        row = conn.execute(
            f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
            (task_id,),
        ).fetchone()
        if not row:
            raise NotFoundError("data process task not found")
        return _decode_row(row) or {}

    @staticmethod
    def _ensure_editable(task: dict[str, Any]) -> None:
        if task["status"] not in EDITABLE_STATUSES:
            raise InvalidStateError(f"task cannot be edited while status is {task['status']}")
        if task.get("output_dataset_id"):
            raise InvalidStateError("published task cannot be edited")

    def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
        allowed = {
            "name",
            "description",
            "process_type",
            "source_dataset_id",
        }
        values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed}
        if payload.get("config") is not None:
            values["config"] = json_dumps(payload["config"])
        if not values:
            return self.get_task(task_id)
        try:
            with self.connect() as conn:
                task = self._task_in_connection(conn, task_id, for_update=True)
                self._ensure_editable(task)
                invalidates_results = (
                    ("config" in payload and payload.get("config") != task.get("config"))
                    or (
                        "process_type" in payload
                        and payload.get("process_type") != task.get("process_type")
                    )
                    or (
                        "source_dataset_id" in payload
                        and payload.get("source_dataset_id") != task.get("source_dataset_id")
                    )
                )
                if invalidates_results:
                    values.update(
                        {
                            "status": "pending",
                            "progress": 0,
                            "output_count": 0,
                            "filtered_count": 0,
                            "duplicate_count": 0,
                            "error_count": 0,
                            "failure_reason": None,
                            "generation_run_id": None,
                            "started_at": None,
                            "completed_at": None,
                        }
                    )
                    conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
                    conn.execute(
                        "DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
                    )
                if (
                    "process_type" in payload
                    and payload.get("process_type") != task.get("process_type")
                ):
                    conn.execute(
                        """
                        UPDATE data_process_source_files
                        SET deleted_at=%s, updated_at=%s
                        WHERE task_id=%s AND deleted_at IS NULL
                        """,
                        (utcnow(), utcnow(), task_id),
                    )
                    values["input_count"] = 0
                values["updated_at"] = utcnow()
                assignments = ", ".join(f"{key}=%s" for key in values)
                row = conn.execute(
                    f"UPDATE data_process_tasks SET {assignments} WHERE id=%s RETURNING *",
                    [*values.values(), task_id],
                ).fetchone()
        except psycopg.errors.UniqueViolation as exc:
            raise ConflictError("data process task name already exists") from exc
        return _decode_row(row) or {}

    def delete_task(self, task_id: str, *, deleted_by: str | None = None) -> None:
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if task["status"] == "running":
                raise InvalidStateError("running task must be stopped before deletion")
            now = utcnow()
            conn.execute(
                """
                UPDATE data_process_tasks
                SET deleted_at=%s, deleted_by=%s, updated_at=%s
                WHERE id=%s
                """,
                (now, deleted_by, now, task_id),
            )

    def list_source_files(self, task_id: str) -> list[dict[str, Any]]:
        self.get_task(task_id)
        with self.connect() as conn:
            rows = conn.execute(
                """
                SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
                       file_format, checksum_sha256, version_no, content_preview, metadata,
                       tenant_id, project_id, created_by, created_at, updated_at
                FROM data_process_source_files
                WHERE task_id=%s AND deleted_at IS NULL
                ORDER BY created_at, id
                """,
                (task_id,),
            ).fetchall()
        return [_decode_row(row) or {} for row in rows]

    def add_source_file(
        self,
        task_id: str,
        *,
        name: str,
        content: str,
        raw_size: int,
        checksum_sha256: str,
        file_format: str,
        record_count: int,
        metadata: dict[str, Any] | None = None,
        created_by: str | None = None,
        source_file_id: str | None = None,
        storage_object_id: str | None = None,
    ) -> dict[str, Any]:
        return self.add_source_files(
            task_id,
            [
                {
                    "name": name,
                    "content": content,
                    "raw_size": raw_size,
                    "checksum_sha256": checksum_sha256,
                    "file_format": file_format,
                    "record_count": record_count,
                    "metadata": metadata or {},
                    "created_by": created_by,
                    "id": source_file_id,
                    "storage_object_id": storage_object_id,
                }
            ],
        )[0]

    def add_source_files(
        self,
        task_id: str,
        files: Sequence[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        """在同一事务中登记一个上传批次，任一文件失败则全部回滚。"""

        if not files:
            raise DataProcessStoreError("at least one source file is required")
        now = utcnow()
        created: list[dict[str, Any]] = []
        try:
            with self.connect() as conn:
                task = self._task_in_connection(conn, task_id, for_update=True)
                self._ensure_editable(task)
                for payload in files:
                    file_id = str(payload.get("id") or new_id("dpsf"))
                    storage_object_id, metadata_payload = _source_storage_descriptor(
                        payload,
                        task_id,
                        file_id,
                    )
                    row = conn.execute(
                        """
                        INSERT INTO data_process_source_files
                        (id, task_id, storage_object_id, name, size_bytes, record_count,
                         file_format, checksum_sha256, version_no, content, content_preview,
                         metadata, tenant_id, project_id,
                         created_by, created_at, updated_at)
                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s, %s,
                                %s, %s, %s, %s)
                        RETURNING id, task_id, storage_object_id, name, size_bytes, record_count,
                                  file_format, checksum_sha256, version_no, content_preview, metadata,
                                  tenant_id, project_id, created_by, created_at, updated_at
                        """,
                        (
                            file_id,
                            task_id,
                            storage_object_id,
                            payload["name"],
                            payload["raw_size"],
                            payload["record_count"],
                            payload["file_format"],
                            payload["checksum_sha256"],
                            payload["content"],
                            str(payload["content"])[:2000],
                            json_dumps(metadata_payload),
                            task.get("tenant_id"),
                            task.get("project_id"),
                            payload.get("created_by") or task.get("created_by"),
                            now,
                            now,
                        ),
                    ).fetchone()
                    created.append(_decode_row(row) or {})
                conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
                conn.execute("DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,))
                conn.execute(
                    """
                    UPDATE data_process_tasks
                    SET status='pending', progress=0, output_count=0, filtered_count=0,
                        duplicate_count=0, error_count=0, failure_reason=NULL,
                        generation_run_id=NULL, started_at=NULL, completed_at=NULL, input_count=(
                        SELECT COALESCE(SUM(record_count), 0)
                        FROM data_process_source_files
                        WHERE task_id=%s AND deleted_at IS NULL
                    ), updated_at=%s
                    WHERE id=%s
                    """,
                    (task_id, now, task_id),
                )
        except psycopg.errors.UniqueViolation as exc:
            raise ConflictError(
                "the same source file content is already attached to this task"
            ) from exc
        return created

    def get_source_file(
        self, task_id: str, file_id: str, *, include_content: bool = True
    ) -> dict[str, Any]:
        # 先验证父任务仍然可见，避免软删除任务后通过已知文件 ID 读取正文。
        self.get_task(task_id)
        content_column = ", content" if include_content else ""
        with self.connect() as conn:
            row = conn.execute(
                f"""
                SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
                       file_format, checksum_sha256, version_no, content_preview, metadata,
                       tenant_id, project_id, created_by, created_at, updated_at{content_column}
                FROM data_process_source_files
                WHERE id=%s AND task_id=%s AND deleted_at IS NULL
                """,
                (file_id, task_id),
            ).fetchone()
        if not row:
            raise NotFoundError("source file not found")
        return _decode_row(row) or {}

    def source_content_window(
        self, task_id: str, file_id: str, offset: int, limit: int
    ) -> dict[str, Any]:
        source_file = self.get_source_file(task_id, file_id, include_content=True)
        content = str(source_file.pop("content", ""))
        window = content[offset : offset + limit]
        return {
            "file": source_file,
            "content": window,
            "offset": offset,
            "limit": limit,
            "total_chars": len(content),
            "has_more": offset + len(window) < len(content),
        }

    def source_content_lines(
        self,
        task_id: str,
        file_id: str,
        start_line: int,
        line_count: int,
    ) -> dict[str, Any]:
        source_file = self.get_source_file(task_id, file_id, include_content=True)
        content = str(source_file.pop("content", ""))
        lines = content.splitlines(keepends=True)
        start_index = min(len(lines), start_line - 1)
        selected = lines[start_index : start_index + line_count]
        end_line = start_index + len(selected)
        return {
            "file": source_file,
            "content": "".join(selected),
            "start_line": start_line,
            "end_line": end_line,
            "line_count": len(selected),
            "total_lines": len(lines),
            "has_more": end_line < len(lines),
        }

    def delete_source_file(self, task_id: str, file_id: str) -> None:
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            self._ensure_editable(task)
            row = conn.execute(
                """
                UPDATE data_process_source_files
                SET deleted_at=%s, updated_at=%s
                WHERE id=%s AND task_id=%s AND deleted_at IS NULL
                RETURNING id
                """,
                (utcnow(), utcnow(), file_id, task_id),
            ).fetchone()
            if not row:
                raise NotFoundError("source file not found")
            conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
            conn.execute(
                "DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_id,)
            )
            conn.execute(
                """
                UPDATE data_process_tasks
                SET status='pending', progress=0, output_count=0, filtered_count=0,
                    duplicate_count=0, error_count=0, failure_reason=NULL,
                    input_count=(SELECT COALESCE(SUM(record_count), 0)
                                 FROM data_process_source_files
                                 WHERE task_id=%s AND deleted_at IS NULL),
                    updated_at=%s
                WHERE id=%s
                """,
                (task_id, utcnow(), task_id),
            )

    def replace_preview_items(
        self,
        task_id: str,
        items: Sequence[dict[str, Any]],
        *,
        source_file_ids: Sequence[str] | None = None,
    ) -> list[dict[str, Any]]:
        selected_ids = (
            list(dict.fromkeys(str(file_id) for file_id in source_file_ids))
            if source_file_ids is not None
            else None
        )
        if selected_ids is not None:
            if not selected_ids or any(not file_id for file_id in selected_ids):
                raise ValueError("source_file_ids must contain non-empty ids")
            selected_set = set(selected_ids)
            unexpected = {
                str(item.get("source_file_id") or "")
                for item in items
                if str(item.get("source_file_id") or "") not in selected_set
            }
            if unexpected:
                raise ValueError("preview items contain an unselected source file")

        now = utcnow()
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            self._ensure_editable(task)
            conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
            if selected_ids is None:
                conn.execute(
                    "DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
                )
            else:
                rows = conn.execute(
                    """
                    SELECT id FROM data_process_source_files
                    WHERE task_id=%s AND deleted_at IS NULL AND id=ANY(%s)
                    """,
                    (task_id, selected_ids),
                ).fetchall()
                found = {str(row["id"]) for row in rows}
                missing = set(selected_ids) - found
                if missing:
                    raise NotFoundError(
                        f"source files not found: {', '.join(sorted(missing))}"
                    )
                conn.execute(
                    """
                    DELETE FROM data_process_preview_items
                    WHERE task_id=%s AND source_file_id=ANY(%s)
                    """,
                    (task_id, selected_ids),
                )
            created: list[dict[str, Any]] = []
            for item in items:
                row = conn.execute(
                    """
                    INSERT INTO data_process_preview_items
                    (id, task_id, source_file_id, original_content, edited_content,
                     source_start, source_end, source_start_line, source_end_line,
                     token_count, status, quality_score, created_at, updated_at)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    RETURNING *
                    """,
                    (
                        item.get("id") or new_id("dpp"),
                        task_id,
                        item.get("source_file_id"),
                        item.get("original_content") or "",
                        item.get("edited_content", item.get("original_content") or ""),
                        item.get("source_start"),
                        item.get("source_end"),
                        item.get("source_start_line"),
                        item.get("source_end_line"),
                        max(0, int(item.get("token_count") or 0)),
                        item.get("status") or "original",
                        json_dumps(item.get("quality_score") or {}),
                        now,
                        now,
                    ),
                ).fetchone()
                created.append(_decode_row(row) or {})
            conn.execute(
                """
                UPDATE data_process_tasks
                SET status='pending', progress=20, output_count=0, filtered_count=0,
                    duplicate_count=0, error_count=0, failure_reason=NULL, updated_at=%s
                WHERE id=%s
                """,
                (now, task_id),
            )
        return created

    def list_preview_items(
        self,
        task_id: str,
        *,
        source_file_id: str | None = None,
        page: int = 1,
        page_size: int = 200,
        keyword: str | None = None,
    ) -> dict[str, Any]:
        self.get_task(task_id)
        clauses = ["task_id=%s"]
        params: list[Any] = [task_id]
        if source_file_id:
            clauses.append("source_file_id=%s")
            params.append(source_file_id)
        if keyword:
            clauses.append("(original_content ILIKE %s OR edited_content ILIKE %s)")
            pattern = f"%{keyword.strip()}%"
            params.extend([pattern, pattern])
        where = " AND ".join(clauses)
        with self.connect() as conn:
            total = conn.execute(
                f"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE {where}", params
            ).fetchone()["count"]
            rows = conn.execute(
                f"""
                SELECT * FROM data_process_preview_items
                WHERE {where}
                ORDER BY source_file_id NULLS LAST, source_start NULLS LAST, created_at, id
                LIMIT %s OFFSET %s
                """,
                [*params, page_size, (page - 1) * page_size],
            ).fetchall()
        return {
            "items": [_decode_row(row) for row in rows],
            "total": int(total),
            "page": page,
            "page_size": page_size,
        }

    def get_preview_item(self, task_id: str, preview_id: str) -> dict[str, Any]:
        with self.connect() as conn:
            row = conn.execute(
                "SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
                (preview_id, task_id),
            ).fetchone()
        if not row:
            raise NotFoundError("preview item not found")
        return _decode_row(row) or {}

    def create_preview_item(self, task_id: str, item: dict[str, Any]) -> dict[str, Any]:
        now = utcnow()
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            self._ensure_editable(task)
            if item.get("source_file_id"):
                source = conn.execute(
                    """
                    SELECT id FROM data_process_source_files
                    WHERE id=%s AND task_id=%s AND deleted_at IS NULL
                    """,
                    (item["source_file_id"], task_id),
                ).fetchone()
                if not source:
                    raise NotFoundError("source file not found")
            row = conn.execute(
                """
                INSERT INTO data_process_preview_items
                (id, task_id, source_file_id, original_content, edited_content,
                 source_start, source_end, source_start_line, source_end_line,
                 token_count, status, quality_score, created_at, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                RETURNING *
                """,
                (
                    new_id("dpp"),
                    task_id,
                    item.get("source_file_id"),
                    item.get("original_content") or "",
                    item.get("edited_content") or "",
                    item.get("source_start"),
                    item.get("source_end"),
                    item.get("source_start_line"),
                    item.get("source_end_line"),
                    max(0, int(item.get("token_count") or 0)),
                    item.get("status") or "manual",
                    json_dumps(item.get("quality_score") or {}),
                    now,
                    now,
                ),
            ).fetchone()
            self._invalidate_results(conn, task_id, now)
        return _decode_row(row) or {}

    def update_preview_item(
        self, task_id: str, preview_id: str, payload: dict[str, Any]
    ) -> dict[str, Any]:
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            self._ensure_editable(task)
            existing = conn.execute(
                "SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
                (preview_id, task_id),
            ).fetchone()
            if not existing:
                raise NotFoundError("preview item not found")
            expected_updated_at = payload.get("expected_updated_at")
            current_updated_at = _serialize_value(existing.get("updated_at"))
            if expected_updated_at and expected_updated_at != current_updated_at:
                raise ConflictError("preview item was modified by another request")
            edited = payload["edited_content"]
            status = payload.get("status")
            if not status:
                if not edited.strip():
                    status = "invalid"
                elif edited == existing["original_content"]:
                    status = "original"
                else:
                    status = "modified"
            now = utcnow()
            row = conn.execute(
                """
                UPDATE data_process_preview_items
                SET edited_content=%s, token_count=%s, status=%s, quality_score=%s,
                    updated_at=%s
                WHERE id=%s AND task_id=%s
                RETURNING *
                """,
                (
                    edited,
                    estimate_token_count(edited),
                    status,
                    json_dumps(payload.get("quality_score") or {}),
                    now,
                    preview_id,
                    task_id,
                ),
            ).fetchone()
            self._invalidate_results(conn, task_id, now)
        return _decode_row(row) or {}

    def delete_preview_item(self, task_id: str, preview_id: str) -> None:
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            self._ensure_editable(task)
            row = conn.execute(
                "DELETE FROM data_process_preview_items WHERE id=%s AND task_id=%s RETURNING id",
                (preview_id, task_id),
            ).fetchone()
            if not row:
                raise NotFoundError("preview item not found")
            self._invalidate_results(conn, task_id, utcnow())

    def _invalidate_results(
        self,
        conn: psycopg.Connection[dict[str, Any]],
        task_id: str,
        now: str,
    ) -> None:
        conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
        conn.execute(
            """
            UPDATE data_process_tasks
            SET status='pending', progress=20, output_count=0, filtered_count=0,
                duplicate_count=0, error_count=0, failure_reason=NULL,
                generation_run_id=NULL, updated_at=%s
            WHERE id=%s
            """,
            (now, task_id),
        )

    def start_generation(self, task_id: str, *, replace_existing: bool = True) -> dict[str, Any]:
        if not replace_existing:
            raise DataProcessStoreError("incremental generation is not supported")
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if task.get("output_dataset_id"):
                raise InvalidStateError("published task cannot be regenerated")
            if task["status"] == "running":
                raise ConflictError("data process task is already running")
            preview_count = conn.execute(
                "SELECT COUNT(*) AS count FROM data_process_preview_items WHERE task_id=%s",
                (task_id,),
            ).fetchone()["count"]
            if not preview_count:
                raise InvalidStateError("preview must be built before generation")
            conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
            now = utcnow()
            generation_run_id = new_id("dprun")
            row = conn.execute(
                """
                UPDATE data_process_tasks
                SET status='running', progress=30, failure_reason=NULL, started_at=%s,
                    completed_at=NULL, filtered_count=0, duplicate_count=0, error_count=0,
                    generation_run_id=%s, updated_at=%s
                WHERE id=%s
                RETURNING *
                """,
                (now, generation_run_id, now, task_id),
            ).fetchone()
        return _decode_row(row) or {}

    def stop_task(self, task_id: str) -> dict[str, Any]:
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if task["status"] != "running":
                raise InvalidStateError("only a running task can be stopped")
            now = utcnow()
            row = conn.execute(
                """
                UPDATE data_process_tasks
                SET status='stopped', failure_reason=NULL, generation_run_id=NULL,
                    updated_at=%s
                WHERE id=%s RETURNING *
                """,
                (now, task_id),
            ).fetchone()
        return _decode_row(row) or {}

    def generation_is_running(self, task_id: str, generation_run_id: str) -> bool:
        task = self.get_task(task_id)
        return (
            task["status"] == "running"
            and task.get("generation_run_id") == generation_run_id
        )

    def update_generation_progress(
        self,
        task_id: str,
        generation_run_id: str,
        processed_count: int,
        total_count: int,
    ) -> bool:
        ratio = processed_count / max(1, total_count)
        progress = min(95.0, 30.0 + ratio * 65.0)
        with self.connect() as conn:
            row = conn.execute(
                """
                UPDATE data_process_tasks
                SET progress=%s, updated_at=%s
                WHERE id=%s AND status='running' AND generation_run_id=%s
                RETURNING id
                """,
                (progress, utcnow(), task_id, generation_run_id),
            ).fetchone()
        return row is not None

    def complete_generation(
        self,
        task_id: str,
        results: Sequence[dict[str, Any]],
        *,
        generation_run_id: str,
        filtered_count: int = 0,
        duplicate_count: int = 0,
        error_count: int = 0,
    ) -> dict[str, Any]:
        now = utcnow()
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if (
                task["status"] != "running"
                or task.get("generation_run_id") != generation_run_id
            ):
                return task
            conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
            for result in results:
                conn.execute(
                    """
                    INSERT INTO data_process_results
                    (id, task_id, preview_item_id, instruction, input, output,
                     original_instruction, original_input, original_output, status, error,
                     split, quality_score, created_at, updated_at)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        result.get("id") or new_id("dpr"),
                        task_id,
                        result.get("preview_item_id"),
                        result.get("instruction") or "",
                        result.get("input") or "",
                        result.get("output") or "",
                        result.get("original_instruction", result.get("instruction") or ""),
                        result.get("original_input", result.get("input") or ""),
                        result.get("original_output", result.get("output") or ""),
                        result.get("status") or "valid",
                        result.get("error"),
                        result.get("split"),
                        json_dumps(result.get("quality_score") or {}),
                        now,
                        now,
                    ),
                )
            row = conn.execute(
                """
                UPDATE data_process_tasks
                SET status='completed', progress=100, output_count=%s, filtered_count=%s,
                    duplicate_count=%s, error_count=%s, failure_reason=NULL,
                    completed_at=%s, generation_run_id=NULL, updated_at=%s
                WHERE id=%s AND generation_run_id=%s RETURNING *
                """,
                (
                    len(results),
                    filtered_count,
                    duplicate_count,
                    error_count,
                    now,
                    now,
                    task_id,
                    generation_run_id,
                ),
            ).fetchone()
        return _decode_row(row) or {}

    def mark_failed(
        self, task_id: str, reason: str, *, generation_run_id: str
    ) -> dict[str, Any]:
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if (
                task["status"] != "running"
                or task.get("generation_run_id") != generation_run_id
            ):
                return task
            now = utcnow()
            row = conn.execute(
                """
                UPDATE data_process_tasks
                SET status='failed', failure_reason=%s, completed_at=%s,
                    generation_run_id=NULL, updated_at=%s
                WHERE id=%s AND generation_run_id=%s RETURNING *
                """,
                (reason[:4000], now, now, task_id, generation_run_id),
            ).fetchone()
        return _decode_row(row) or {}

    def progress(self, task_id: str) -> dict[str, Any]:
        task = self.get_task(task_id)
        return {
            "task_id": task["id"],
            "status": task["status"],
            "progress": float(task.get("progress") or 0),
            "input_count": int(task.get("input_count") or 0),
            "output_count": int(task.get("output_count") or 0),
            "filtered_count": int(task.get("filtered_count") or 0),
            "duplicate_count": int(task.get("duplicate_count") or 0),
            "error_count": int(task.get("error_count") or 0),
            "failure_reason": task.get("failure_reason"),
            "started_at": task.get("started_at"),
            "completed_at": task.get("completed_at"),
        }

    def list_results(
        self,
        task_id: str,
        *,
        page: int = 1,
        page_size: int = 100,
        status: str | None = None,
        split: str | None = None,
        keyword: str | None = None,
    ) -> dict[str, Any]:
        self.get_task(task_id)
        clauses = ["task_id=%s"]
        params: list[Any] = [task_id]
        if status:
            clauses.append("status=%s")
            params.append(status)
        if split:
            clauses.append("split=%s")
            params.append(split)
        if keyword:
            clauses.append("(instruction ILIKE %s OR input ILIKE %s OR output ILIKE %s)")
            pattern = f"%{keyword.strip()}%"
            params.extend([pattern, pattern, pattern])
        where = " AND ".join(clauses)
        with self.connect() as conn:
            total = conn.execute(
                f"SELECT COUNT(*) AS count FROM data_process_results WHERE {where}", params
            ).fetchone()["count"]
            rows = conn.execute(
                f"""
                SELECT * FROM data_process_results WHERE {where}
                ORDER BY created_at, id LIMIT %s OFFSET %s
                """,
                [*params, page_size, (page - 1) * page_size],
            ).fetchall()
        return {
            "items": [_decode_row(row) for row in rows],
            "total": int(total),
            "page": page,
            "page_size": page_size,
        }

    def get_result(self, task_id: str, result_id: str) -> dict[str, Any]:
        with self.connect() as conn:
            row = conn.execute(
                "SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
                (result_id, task_id),
            ).fetchone()
        if not row:
            raise NotFoundError("data process result not found")
        return _decode_row(row) or {}

    def update_result(
        self, task_id: str, result_id: str, payload: dict[str, Any]
    ) -> dict[str, Any]:
        allowed = {"instruction", "input", "output", "quality_score"}
        values = {key: value for key, value in payload.items() if key in allowed}
        if "quality_score" in values:
            values["quality_score"] = json_dumps(values["quality_score"])
        if not values:
            raise DataProcessStoreError("no result fields supplied")
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if task["status"] == "running":
                raise InvalidStateError("results cannot be edited while generation is running")
            if task.get("output_dataset_id"):
                raise InvalidStateError("published results cannot be edited")
            current = conn.execute(
                "SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
                (result_id, task_id),
            ).fetchone()
            if not current:
                raise NotFoundError("data process result not found")
            expected_updated_at = payload.get("expected_updated_at")
            current_updated_at = _serialize_value(current.get("updated_at"))
            if expected_updated_at and expected_updated_at != current_updated_at:
                raise ConflictError("data process result was modified by another request")
            merged = {**current, **values}
            quality = payload.get("quality_score") or {}
            hard_valid = bool(
                str(merged.get("instruction") or "").strip()
                and str(merged.get("output") or "").strip()
            )
            quality_valid = bool(quality.get("is_valid", hard_valid))
            changed = any(
                str(merged.get(field) or "")
                != str(merged.get(f"original_{field}") or "")
                for field in ("instruction", "input", "output")
            )
            status = "invalid" if not hard_valid or not quality_valid else (
                "modified" if changed else "valid"
            )
            values["status"] = status
            flags = quality.get("flags") if isinstance(quality, dict) else None
            values["error"] = ", ".join(str(flag) for flag in flags or []) or (
                "quality validation failed" if status == "invalid" else None
            )
            values["updated_at"] = utcnow()
            assignments = ", ".join(f"{key}=%s" for key in values)
            row = conn.execute(
                f"""UPDATE data_process_results SET {assignments}
                    WHERE id=%s AND task_id=%s RETURNING *""",
                [*values.values(), result_id, task_id],
            ).fetchone()
            conn.execute(
                """
                UPDATE data_process_tasks
                SET error_count=(
                    SELECT COUNT(*) FROM data_process_results
                    WHERE task_id=%s AND status='invalid'
                ), updated_at=%s
                WHERE id=%s
                """,
                (task_id, utcnow(), task_id),
            )
        return _decode_row(row) or {}

    def get_generation_model(self, model_id: str) -> dict[str, Any]:
        with self.connect() as conn:
            row = conn.execute(
                """
                SELECT id, name, type, purpose, model_source, description, path,
                       api_url, api_key, online_model_name, create_time
                FROM models WHERE id=%s
                """,
                (model_id,),
            ).fetchone()
        if not row:
            raise NotFoundError("generation model not found")
        return _decode_row(row) or {}

    def save_generation_model_snapshot(
        self,
        task_id: str,
        model_snapshot: dict[str, Any],
        *,
        generation_run_id: str,
    ) -> dict[str, Any]:
        # API 密钥仅用于本次调用，绝不能进入任务配置、详情响应或审计快照。
        safe_snapshot = {
            key: value for key, value in model_snapshot.items() if key != "api_key"
        }
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if (
                task["status"] != "running"
                or task.get("generation_run_id") != generation_run_id
            ):
                raise InvalidStateError("generation run is no longer active")
            config = dict(task.get("config") or {})
            config["generation_model_snapshot"] = safe_snapshot
            row = conn.execute(
                """
                UPDATE data_process_tasks SET config=%s, updated_at=%s
                WHERE id=%s AND generation_run_id=%s RETURNING *
                """,
                (json_dumps(config), utcnow(), task_id, generation_run_id),
            ).fetchone()
        return _decode_row(row) or {}

    def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
        """发布有效结果；任务行锁保证重复请求返回同一数据集。"""
        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if task.get("output_dataset_id"):
                dataset = conn.execute(
                    "SELECT * FROM datasets WHERE id=%s", (task["output_dataset_id"],)
                ).fetchone()
                if dataset:
                    return {"dataset": _decode_row(dataset), "created": False}
                # 数据集被外部流程清理后，解除断链并重新发布。
                conn.execute(
                    "UPDATE data_process_tasks SET output_dataset_id=NULL WHERE id=%s",
                    (task_id,),
                )
            if task["status"] != "completed":
                raise InvalidStateError("only a completed task can be published")
            rows = conn.execute(
                """
                SELECT * FROM data_process_results
                WHERE task_id=%s ORDER BY created_at, id
                """,
                (task_id,),
            ).fetchall()
            if not rows:
                raise InvalidStateError("task has no results to publish")
            invalid_count = sum(
                1
                for row in rows
                if row["status"] == "invalid"
                or not str(row.get("instruction") or "").strip()
                or not str(row.get("output") or "").strip()
            )
            if invalid_count:
                raise InvalidStateError(f"task contains {invalid_count} invalid results")

            dataset_id = new_id("dataset")
            file_id = new_id("dfile")
            version_id = new_id("dfv")
            now = utcnow()
            requested_split = payload.get("split") or {
                "train": 80,
                "validation": 10,
                "test": 10,
            }
            records = [
                {
                    "instruction": row["instruction"],
                    "input": row["input"],
                    "output": row["output"],
                    "split": stable_split(
                        str(row["id"]),
                        requested_split,
                        seed=task_id,
                    ),
                }
                for row in rows
            ]
            content = "".join(json_dumps(record) + "\n" for record in records)
            raw = content.encode("utf-8")
            checksum = hashlib.sha256(raw).hexdigest()
            storage_object_id = f"db://data-process/{task_id}/{file_id}/v1"
            source_result_ids = [row["id"] for row in rows]
            metadata = {
                "source": "data_process",
                "storage_backend": "database",
                "storage_object_id": storage_object_id,
                "source_task_id": task_id,
                "source_file_ids": [item["id"] for item in self._source_ids(conn, task_id)],
                "source_result_ids": source_result_ids,
                "format": payload.get("format") or "alpaca_jsonl",
                "split": payload.get("split") or {},
            }
            try:
                dataset = conn.execute(
                    """
                    INSERT INTO datasets
                    (id, name, type, storage_type, source, task_id, source_task_id,
                     size, size_bytes, count, record_count, description, metadata,
                     tenant_id, project_id, owner_id, created_by, create_time,
                     created_at, updated_at)
                    VALUES (%s, %s, %s, %s, 'task', %s, %s, %s, %s, %s, %s, %s, %s,
                            %s, %s, %s, %s, %s, %s, %s)
                    RETURNING *
                    """,
                    (
                        dataset_id,
                        payload["dataset_name"],
                        payload.get("dataset_type") or "train",
                        payload.get("storage_type") or "local",
                        task_id,
                        task_id,
                        f"{len(raw)} B",
                        len(raw),
                        len(records),
                        len(records),
                        payload.get("description") or task.get("description") or "",
                        json_dumps(metadata),
                        task.get("tenant_id"),
                        task.get("project_id"),
                        task.get("owner_id"),
                        payload.get("created_by") or task.get("created_by"),
                        now,
                        now,
                        now,
                    ),
                ).fetchone()
                version = {
                    "id": version_id,
                    "version_no": 1,
                    "description": "data process publish",
                    "checksum_sha256": checksum,
                    "size_bytes": len(raw),
                    "record_count": len(records),
                    "created_at": now,
                    "source_task_id": task_id,
                    "storage_object_id": storage_object_id,
                }
                conn.execute(
                    """
                    INSERT INTO dataset_files
                    (id, dataset_id, name, storage_object_id, size, content,
                     active_version_id, versions, create_time, current_version_id,
                     size_bytes, record_count, file_format, checksum_sha256, version_no,
                     source_task_id, tenant_id, project_id, created_by, metadata,
                     created_at, updated_at)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
                            %s, %s, 1, %s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        file_id,
                        dataset_id,
                        f"{payload['dataset_name']}.jsonl",
                        storage_object_id,
                        f"{len(raw)} B",
                        content,
                        version_id,
                        json_dumps([version]),
                        now,
                        version_id,
                        len(raw),
                        len(records),
                        "jsonl",
                        checksum,
                        task_id,
                        task.get("tenant_id"),
                        task.get("project_id"),
                        payload.get("created_by") or task.get("created_by"),
                        json_dumps(metadata),
                        now,
                        now,
                    ),
                )
                conn.execute(
                    """
                    INSERT INTO dataset_file_versions
                    (id, dataset_file_id, version_no, storage_object_id, content_preview,
                     description, size_bytes, record_count, checksum_sha256,
                     source_task_id, metadata, created_by, created_at)
                    VALUES (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        version_id,
                        file_id,
                        storage_object_id,
                        content[:2000],
                        "data process publish",
                        len(raw),
                        len(records),
                        checksum,
                        task_id,
                        json_dumps(metadata),
                        payload.get("created_by") or task.get("created_by"),
                        now,
                    ),
                )
                for line_number, (source_row, record) in enumerate(
                    zip(rows, records, strict=True), start=1
                ):
                    conn.execute(
                        """
                        INSERT INTO dataset_records
                        (id, dataset_id, dataset_file_id, version_id, line_no, split,
                         instruction, input, output, raw, status, source_task_id,
                         source_result_id, preview_item_id, created_at)
                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
                                %s, %s, %s, %s)
                        """,
                        (
                            new_id("drec"),
                            dataset_id,
                            file_id,
                            version_id,
                            line_number,
                            record["split"],
                            record["instruction"],
                            record["input"],
                            record["output"],
                            json_dumps(
                                {
                                    **record,
                                    "source_task_id": task_id,
                                    "source_result_id": source_row["id"],
                                    "preview_item_id": source_row.get("preview_item_id"),
                                }
                            ),
                            source_row["status"],
                            task_id,
                            source_row["id"],
                            source_row.get("preview_item_id"),
                            now,
                        ),
                    )
            except psycopg.errors.UniqueViolation as exc:
                raise ConflictError("dataset name already exists") from exc
            conn.execute(
                """
                UPDATE data_process_tasks
                SET output_dataset_id=%s, updated_at=%s, updated_by=%s
                WHERE id=%s
                """,
                (dataset_id, now, payload.get("created_by"), task_id),
            )
        return {"dataset": _decode_row(dataset), "created": True}

    @staticmethod
    def _source_ids(
        conn: psycopg.Connection[dict[str, Any]], task_id: str
    ) -> list[dict[str, Any]]:
        return conn.execute(
            """
            SELECT id FROM data_process_source_files
            WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id
            """,
            (task_id,),
        ).fetchall()


@lru_cache
def get_data_process_store() -> DataProcessStore:
    return DataProcessStore()