"""数据处理存储层 - 结果管理。"""

from __future__ import annotations

from typing import Any
from collections.abc import Sequence

import psycopg

from .base import (
    StoreBase,
    utcnow,
    new_id,
    repeat_task_id,
    json_dumps,
    _json_value,
    _decode_row,
    _public_task,
    _business_config,
    _preview_config_value,
    _preview_config_changed,
    _preview_config_projection,
    _normalized_preprocess_options,
    _regeneration_marker,
    _is_regeneration_prepared,
    _task_output_type,
    _task_reasoning_detail,
    _reasoning_output_is_valid,
    _dpo_fields_are_valid,
    _source_storage_descriptor,
    _serialize_value,
    DataProcessStoreError,
    NotFoundError,
    ConflictError,
    InvalidStateError,
    EDITABLE_STATUSES,
    ACTIVE_PREVIEW_STATUSES,
    WORKFLOW_STEPS,
    _REGENERATION_MARKER_KEY,
    _REPEAT_SOURCE_TASK_KEY,
    _REPEAT_REQUEST_KEY,
)


class ResultsMixin:
    """结果管理 Mixin。"""

    def confirm_results(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"] != "completed":
                raise InvalidStateError("only a completed task can confirm results")
            if task.get("workflow_step") != "results":
                raise InvalidStateError("workflow must be on results before confirmation")
            if task.get("results_confirmed"):
                return task

            rows = conn.execute(
                """
                SELECT status, instruction, output, chosen, rejected
                FROM data_process_results
                WHERE task_id=%s
                """,
                (task_id,),
            ).fetchall()
            if not rows:
                raise InvalidStateError("task has no results to confirm")
            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()
                or (
                    _task_output_type(task) == "reasoning"
                    and not _reasoning_output_is_valid(row.get("output"))
                )
                or (
                    _task_output_type(task) == "dpo"
                    and not _dpo_fields_are_valid(row)
                )
            )
            if invalid_count:
                raise InvalidStateError(
                    f"task contains {invalid_count} invalid results"
                )

            row = conn.execute(
                """
                UPDATE data_process_tasks
                SET results_confirmed=TRUE, updated_at=%s
                WHERE id=%s RETURNING *
                """,
                (utcnow(), task_id),
            ).fetchone()
        return _decode_row(row) or {}

    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", "chosen", "rejected", "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")
            output_type = _task_output_type(task)
            if output_type == "dpo" and "chosen" in values:
                values["output"] = values["chosen"]
            merged = {**current, **values}
            quality = payload.get("quality_score") or {}
            instruction_valid = bool(str(merged.get("instruction") or "").strip())
            output_valid = bool(str(merged.get("output") or "").strip())
            reasoning_valid = (
                output_type != "reasoning"
                or _reasoning_output_is_valid(merged.get("output"))
            )
            dpo_valid = output_type != "dpo" or _dpo_fields_are_valid(merged)
            hard_valid = instruction_valid and output_valid and reasoning_valid and dpo_valid
            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", "chosen", "rejected")
                    if output_type == "dpo"
                    else ("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
            format_error = (
                "思维链输出必须包含非空的 <think>...</think> 推理过程和最终答案"
                if instruction_valid and output_valid and not reasoning_valid
                else "DPO 输出必须包含不同的非空 Chosen 和 Rejected 回答"
                if instruction_valid and not dpo_valid
                else "Instruction 和 Output 不能为空"
                if not instruction_valid or not output_valid
                else None
            )
            values["error"] = ", ".join(str(flag) for flag in flags or []) or (
                format_error
                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 replace_generated_result(
        self,
        task_id: str,
        result_id: str,
        replacement: dict[str, Any],
        *,
        expected_updated_at: str,
    ) -> dict[str, Any]:
        """用新模型结果原位替换失败项，并将新内容设为恢复基线。"""

        with self.connect() as conn:
            task = self._task_in_connection(conn, task_id, for_update=True)
            if task["status"] != "completed" or task.get("workflow_step") != "results":
                raise InvalidStateError("task is not editing generation results")
            if task.get("results_confirmed"):
                raise InvalidStateError("confirmed results cannot be regenerated")
            if task.get("output_dataset_id"):
                raise InvalidStateError("published results cannot be regenerated")

            current = conn.execute(
                """SELECT * FROM data_process_results
                   WHERE id=%s AND task_id=%s FOR UPDATE""",
                (result_id, task_id),
            ).fetchone()
            if not current:
                raise NotFoundError("data process result not found")
            if current.get("status") != "invalid":
                raise InvalidStateError("only an invalid result can be regenerated")
            current_updated_at = _serialize_value(current.get("updated_at"))
            if expected_updated_at != current_updated_at:
                raise ConflictError("data process result was modified by another request")

            instruction = str(replacement.get("instruction") or "").strip()
            input_text = str(replacement.get("input") or "").strip()
            output = str(replacement.get("output") or "").strip()
            chosen = str(replacement.get("chosen") or "").strip()
            rejected = str(replacement.get("rejected") or "").strip()
            quality_score = replacement.get("quality_score") or {}
            if not instruction or not output or not bool(quality_score.get("is_valid")):
                raise InvalidStateError("regenerated result did not pass quality validation")
            if _task_output_type(task) == "reasoning" and not _reasoning_output_is_valid(output):
                raise InvalidStateError("regenerated reasoning result has an invalid output format")
            if _task_output_type(task) == "dpo" and not _dpo_fields_are_valid(replacement):
                raise InvalidStateError("regenerated DPO result has invalid preference fields")

            now = utcnow()
            row = conn.execute(
                """
                UPDATE data_process_results
                SET instruction=%s, input=%s, output=%s, chosen=%s, rejected=%s,
                    original_instruction=%s, original_input=%s, original_output=%s,
                    original_chosen=%s, original_rejected=%s,
                    status='valid', error=NULL, quality_score=%s, updated_at=%s
                WHERE id=%s AND task_id=%s
                RETURNING *
                """,
                (
                    instruction,
                    input_text,
                    output,
                    chosen,
                    rejected,
                    instruction,
                    input_text,
                    output,
                    chosen,
                    rejected,
                    json_dumps(quality_score),
                    now,
                    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, now, task_id),
            )
        return _decode_row(row) or {}