"""Transactional exploration records in the existing automation database.

The queue remains the sole scheduler. These records describe a conversation and
its immutable evidence/case versions, not a second execution state machine.
"""

from __future__ import annotations

from collections.abc import Callable
from contextlib import contextmanager
import hashlib
import json
import time
from typing import Literal
from uuid import uuid4

from apps.admin_console.services.automation_store import AutomationStore, identifier, store
from artemis.config.exploration import ExplorationConfig
from artemis.exploration.models import CaseDraft, ExplorationStatus, RunKind


class ExplorationConflict(ValueError):
    """A stale request cannot overwrite newer conversation or run state."""


def _digest(value: object) -> str:
    return hashlib.sha256(
        json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()


class ExplorationStore:
    def __init__(self, database: AutomationStore | None = None):
        self.database = database if database is not None else store

    @contextmanager
    def _transaction(self):
        with self.database.connection() as connection:
            connection.execute("BEGIN IMMEDIATE")
            try:
                yield connection
                connection.commit()
            except BaseException:
                connection.rollback()
                raise

    @staticmethod
    def _get(connection, kind: str, key: str) -> dict | None:
        row = connection.execute(
            "SELECT payload FROM automation_records WHERE kind=? AND id=?", (kind, key)
        ).fetchone()
        return json.loads(row[0]) if row else None

    @staticmethod
    def _put(connection, kind: str, key: str, payload: dict) -> dict:
        connection.execute(
            "INSERT INTO automation_records VALUES (?,?,?,?) ON CONFLICT(kind,id) "
            "DO UPDATE SET payload=excluded.payload,updated_at=excluded.updated_at",
            (kind, key, json.dumps(payload, ensure_ascii=False), time.time()),
        )
        return payload

    def _require(self, connection, session_id: str) -> dict:
        session = self._get(connection, "exploration", identifier(session_id))
        if session is None:
            raise KeyError("探索会话不存在")
        if session.get("status") == "deleting":
            raise ExplorationConflict("探索正在删除，不能继续操作")
        return session

    def _operation(
        self, connection, scope: str, request_id: str, value: object, operation: Callable[[], dict]
    ) -> dict:
        if not request_id.strip() or len(request_id) > 200:
            raise ValueError("request_id 必须在 1–200 字之间")
        key = _digest([scope, request_id])
        previous = self._get(connection, "exploration_request", key)
        fingerprint = _digest(value)
        if previous:
            if previous["fingerprint"] != fingerprint:
                raise ExplorationConflict("同一 request_id 的内容与已提交请求不同")
            if previous.get("deleted"):
                raise ExplorationConflict("此请求关联的探索已删除，请重新创建任务")
            return previous["result"]
        result = operation()
        self._put(
            connection,
            "exploration_request",
            key,
            {
                "fingerprint": fingerprint,
                "result": result,
            },
        )
        return result

    def _save_session(self, connection, session: dict) -> None:
        session["updated_at"] = time.time()
        self._put(connection, "exploration", session["id"], session)

    def _append_message(
        self, connection, session: dict, role: str, text: str, source: str, **metadata
    ) -> dict:
        sequence = session["message_count"] + 1
        message = {
            "id": str(uuid4()),
            "exploration_id": session["id"],
            "sequence": sequence,
            "role": role,
            "text": text,
            "source": source,
            "created_at": time.time(),
            **metadata,
        }
        self._put(connection, "exploration_message", f"{session['id']}_{sequence:010d}", message)
        session["message_count"] = sequence
        self._save_session(connection, session)
        return message

    def create(
        self,
        goal: str,
        config: ExplorationConfig,
        task_config: dict,
        request_id: str,
        *,
        source: str = "web",
    ) -> dict:
        if not goal.strip() or len(goal) > 10000:
            raise ValueError("探索目标必须在 1–10000 字之间")
        value = {"goal": goal.strip(), "config": config.model_dump(), "task_config": task_config}
        with self._transaction() as connection:

            def create():
                session = {
                    "id": str(uuid4()),
                    **value,
                    "status": "draft",
                    "message_count": 0,
                    "case_version": 0,
                    "active_run_id": None,
                    "pending_question": None,
                    "active_demonstration_id": None,
                    "created_at": time.time(),
                    "updated_at": time.time(),
                    "elapsed_seconds": 0.0,
                    "continuation_count": 0,
                }
                self._append_message(connection, session, "user", value["goal"], source)
                return session

            return self._operation(connection, "create", request_id, value, create)

    def get(self, session_id: str) -> dict:
        with self.database.connection() as connection:
            return self._require(connection, session_id)

    def list(self, limit: int = 100) -> list[dict]:
        return self.database.records("exploration", min(max(limit, 1), 500))

    def messages(self, session_id: str, *, after: int = 0, limit: int = 100) -> list[dict]:
        self.get(session_id)
        with self.database.connection() as connection:
            rows = connection.execute(
                "SELECT payload FROM automation_records WHERE kind='exploration_message' "
                "AND id>=? AND id<=? ORDER BY id LIMIT ?",
                (
                    f"{session_id}_{max(after + 1, 1):010d}",
                    f"{session_id}_9999999999",
                    min(max(limit, 1), 100),
                ),
            ).fetchall()
        return [json.loads(row[0]) for row in rows]

    def add_message(
        self,
        session_id: str,
        text: str,
        request_id: str,
        *,
        source: str = "web",
        role: Literal["user", "assistant", "system"] = "user",
        metadata: dict | None = None,
    ) -> dict:
        if not text.strip() or len(text) > 10000:
            raise ValueError("消息必须在 1–10000 字之间")
        metadata = metadata or {}
        if set(metadata) - {
            "run_id",
            "question_id",
            "demonstration_id",
            "kind",
            "recommendation",
            "reply_to",
        }:
            raise ValueError("未知的探索消息元数据")
        with self._transaction() as connection:
            session = self._require(connection, session_id)

            def append():
                run_id = metadata.get("run_id")
                run = self._get(connection, "exploration_run", run_id) if run_id else None
                if role == "user" and run and run.get("status") == "finalizing":
                    raise ExplorationConflict(
                        "本轮操作已结束，正在提炼用例。请稍后继续探索；指导尚未发送。"
                    )
                message = self._append_message(
                    connection, session, role, text.strip(), source, **metadata
                )
                if role == "user" and run and run.get("outcome") is None:
                    run["user_guidance_count"] = run.get("user_guidance_count", 0) + 1
                    self._put(connection, "exploration_run", run_id, run)
                return message

            return self._operation(
                connection,
                f"{session_id}:message",
                request_id,
                {
                    "text": text.strip(),
                    "role": role,
                    "metadata": metadata if role != "user" else {},
                },
                append,
            )

    def recent_messages(
        self, session_id: str, *, characters: int = 24000, guidance_only: bool = False
    ) -> list[dict]:
        """Bound model context without deleting the durable conversation."""
        self.get(session_id)
        with self.database.connection() as connection:
            rows = connection.execute(
                "SELECT payload FROM automation_records WHERE kind='exploration_message' "
                "AND id>=? AND id<=? AND (?=0 OR json_extract(payload,'$.role')!='assistant' "
                "OR COALESCE(json_extract(payload,'$.kind'),'') IN ('question','demonstration')) "
                "ORDER BY id DESC LIMIT 100",
                (f"{session_id}_0000000001", f"{session_id}_9999999999", int(guidance_only)),
            ).fetchall()
        result = []
        used = 0
        for row in rows:
            message = json.loads(row[0])
            size = len(message["text"]) + len(message.get("recommendation", ""))
            if result and used + size > characters:
                break
            result.append(message)
            used += size
        return list(reversed(result))

    def reserve_run(
        self,
        session_id: str,
        request_id: str,
        *,
        kind: RunKind,
        goal: str,
        task_config: dict,
        case_version: int | None = None,
    ) -> dict:
        value = {
            "kind": kind,
            "goal": goal,
            "task_config": task_config,
            "case_version": case_version,
        }
        with self._transaction() as connection:
            session = self._require(connection, session_id)

            def reserve():
                if session["active_run_id"]:
                    raise ExplorationConflict("该探索已有运行中的任务，请先完成或停止它")
                if kind == "validation":
                    if not case_version or not self._get(
                        connection, "exploration_case", f"{session_id}_{case_version:010d}"
                    ):
                        raise ValueError("复跑必须指定已保存的用例版本")
                elif case_version is not None:
                    raise ValueError("探索运行不能冒充用例复跑")
                run = {
                    "id": str(uuid4()),
                    "exploration_id": session_id,
                    **value,
                    "status": "reserved",
                    "created_at": time.time(),
                    "outcome": None,
                }
                self._put(connection, "exploration_run", run["id"], run)
                session.update(active_run_id=run["id"], status="queued", pending_question=None)
                self._save_session(connection, session)
                return run

            return self._operation(connection, f"{session_id}:run", request_id, value, reserve)

    def run(self, run_id: str) -> dict | None:
        return self.database.get("exploration_run", identifier(run_id))

    def runs(self, session_id: str) -> list[dict]:
        self.get(session_id)
        with self.database.connection() as connection:
            rows = connection.execute(
                "SELECT payload FROM automation_records WHERE kind='exploration_run' "
                "AND json_extract(payload,'$.exploration_id')=? ORDER BY updated_at",
                (session_id,),
            ).fetchall()
        return [json.loads(row[0]) for row in rows]

    def begin_delete(self, session_id: str, check_runs: Callable[[list[dict]], None]) -> list[dict]:
        """Fence new work before deleting files; failed cleanup can be retried."""
        with self._transaction() as connection:
            session = self._get(connection, "exploration", identifier(session_id))
            if session is None:
                raise KeyError("探索会话不存在")
            if session.get("active_run_id") or session.get("active_demonstration_id"):
                raise ExplorationConflict("探索仍有执行或示范任务，请先停止并等待结束后删除")
            rows = connection.execute(
                "SELECT kind,payload FROM automation_records "
                "WHERE json_extract(payload,'$.exploration_id')=?",
                (session_id,),
            ).fetchall()
            runs = []
            for row in rows:
                item = json.loads(row["payload"])
                if row["kind"] == "exploration_run":
                    if item.get("outcome") is None:
                        raise ExplorationConflict("探索执行尚未结束，请先停止并等待结束后删除")
                    runs.append(item)
                elif row["kind"] in {"exploration_generation", "exploration_discussion"}:
                    if item.get("status") in {"pending", "running"}:
                        raise ExplorationConflict("探索正在分析或生成内容，请等待结束后删除")
            check_runs(runs)
            session["status"] = "deleting"
            self._save_session(connection, session)
            return runs

    def finish_delete(self, session_id: str) -> None:
        with self._transaction() as connection:
            session = self._get(connection, "exploration", session_id)
            if not session or session.get("status") != "deleting":
                raise ExplorationConflict("探索尚未准备删除")
            connection.execute(
                "DELETE FROM automation_records WHERE kind IN "
                "('exploration_run','exploration_message','exploration_case',"
                "'exploration_demo','exploration_generation','exploration_discussion') "
                "AND json_extract(payload,'$.exploration_id')=?",
                (session_id,),
            )
            connection.execute(
                "DELETE FROM automation_records WHERE kind='exploration' AND id=?", (session_id,)
            )
            # Retain only idempotency tombstones, not snapshots of deleted messages/cases.
            connection.execute(
                "UPDATE automation_records SET payload="
                "json_set(json_remove(payload,'$.result'),'$.deleted',json('true')) "
                "WHERE kind='exploration_request' AND "
                "(json_extract(payload,'$.result.id')=? OR "
                "json_extract(payload,'$.result.exploration_id')=?)",
                (session_id, session_id),
            )

    def set_run_state(
        self,
        run_id: str,
        status: ExplorationStatus,
        *,
        reason: str | None = None,
        elapsed_seconds: float | None = None,
        expected_status: str | None = None,
    ) -> dict:
        if status not in {
            "queued",
            "running",
            "waiting_user",
            "demonstrating",
            "analyzing_demo",
            "paused",
        }:
            raise ValueError("执行结束必须保存真实结果，不能直接改状态")
        with self._transaction() as connection:
            run = self._get(connection, "exploration_run", identifier(run_id))
            if run is None:
                raise KeyError("探索运行不存在")
            session = self._require(connection, run["exploration_id"])
            if expected_status is not None and run["status"] != expected_status:
                return session
            if session["active_run_id"] != run_id or run.get("outcome") is not None:
                raise ExplorationConflict("该执行已结束，不能覆盖当前探索状态")
            if elapsed_seconds is not None:
                previous = float(run.get("elapsed_seconds", 0))
                run["elapsed_seconds"] = max(previous, elapsed_seconds)
                session["elapsed_seconds"] += run["elapsed_seconds"] - previous
            run.update(status=status, reason=reason)
            self._put(connection, "exploration_run", run_id, run)
            session.update(status=status, pause_reason=reason)
            self._save_session(connection, session)
            return session

    def finish_run(self, run_id: str, outcome: dict) -> dict:
        """Retain actual verdicts; a completed run is not automatically a pass."""
        with self._transaction() as connection:
            run = self._get(connection, "exploration_run", identifier(run_id))
            if run is None:
                raise KeyError("探索运行不存在")
            session = self._require(connection, run["exploration_id"])
            if run.get("outcome") is not None:
                return run
            # Process-exit metadata may contain only a lifecycle status. Keep
            # the checker's recorded evidence even when the worker exits normally.
            outcome = {
                **(run.get("verified_outcome") or {}),
                **{key: value for key, value in outcome.items() if value is not None},
            }
            run.update(
                outcome=outcome,
                status=outcome.get("status", "interrupted"),
                finished_at=time.time(),
            )
            self._put(connection, "exploration_run", run_id, run)
            if session["active_run_id"] == run_id:
                if demo_id := session.get("active_demonstration_id"):
                    demo = self._get(connection, "exploration_demo", demo_id)
                    if demo and demo.get("status") not in {
                        "completed",
                        "failed",
                        "cancelled",
                        "timed_out",
                    }:
                        demo.update(
                            status="failed",
                            error="执行已结束，示范被中断；已写入的文件保留在任务资料中",
                            finished_at=time.time(),
                        )
                        self._put(connection, "exploration_demo", demo_id, demo)
                session.update(
                    active_run_id=None,
                    active_demonstration_id=None,
                    status="completed"
                    if run["status"] == "completed"
                    else ("cancelled" if run["status"] == "cancelled" else "paused"),
                    pending_question=None,
                )
                self._save_session(connection, session)
            from artemis.runtime.case_result import resolve_case_result

            result = resolve_case_result(outcome, run["status"])
            label = "复跑" if run["kind"] == "validation" else "探索"
            self._append_message(
                connection,
                session,
                "system",
                f"{label}执行已结束（{run['status']}）。"
                + (f"验收：{result.reason}" if result else "验收尚未确定。"),
                "exploration",
                run_id=run_id,
                kind="outcome",
            )
            return run

    def record_verdict(self, run_id: str, outcome: dict) -> None:
        """Save checker evidence while synthesis still owns the live execution lease."""
        with self._transaction() as connection:
            run = self._get(connection, "exploration_run", identifier(run_id))
            if not run or run["kind"] != "exploration" or run.get("outcome") is not None:
                raise ExplorationConflict("探索执行已结束")
            run["verified_outcome"] = outcome
            self._put(connection, "exploration_run", run_id, run)

    def seal_execution(self, run_id: str, guidance_count: int) -> bool:
        """Close admission atomically with the last reviewed guidance revision."""
        with self._transaction() as connection:
            run = self._get(connection, "exploration_run", run_id)
            if not run or run.get("outcome") is not None:
                return False
            session = self._require(connection, run["exploration_id"])
            if (
                session["active_run_id"] != run_id
                or run.get("user_guidance_count", 0) != guidance_count
                or session.get("pending_question")
                or session.get("active_demonstration_id")
                or session["status"] in {"paused", "waiting_user", "demonstrating"}
            ):
                return False
            run["status"] = "finalizing"
            session["status"] = "finalizing"
            self._put(connection, "exploration_run", run_id, run)
            self._save_session(connection, session)
            return True

    def record_elapsed(self, run_id: str, elapsed_seconds: float, epoch: int = 0) -> None:
        with self._transaction() as connection:
            run = self._get(connection, "exploration_run", run_id)
            if not run or run.get("outcome") is not None:
                return
            session = self._require(connection, run["exploration_id"])
            previous = float(run.get("elapsed_seconds", 0))
            run["elapsed_seconds"] = max(previous, elapsed_seconds)
            session["elapsed_seconds"] += run["elapsed_seconds"] - previous
            previous_epoch = int(run.get("budget_epoch", 0))
            if epoch > previous_epoch:
                session["continuation_count"] += epoch - previous_epoch
                run["budget_epoch"] = epoch
            self._put(connection, "exploration_run", run_id, run)
            self._save_session(connection, session)

    def ask(self, run_id: str, question: str, recommendation: str, request_id: str) -> dict:
        with self._transaction() as connection:
            run = self._get(connection, "exploration_run", identifier(run_id))
            if not run:
                raise KeyError("探索运行不存在")
            session = self._require(connection, run["exploration_id"])

            def ask():
                if session["active_run_id"] != run_id:
                    raise ExplorationConflict("只能向当前探索提出问题")
                message = self._append_message(
                    connection,
                    session,
                    "assistant",
                    question,
                    "exploration",
                    run_id=run_id,
                    kind="question",
                    recommendation=recommendation,
                )
                session.update(pending_question=message["id"], status="waiting_user")
                self._save_session(connection, session)
                return message

            return self._operation(
                connection, f"{run_id}:question", request_id, [question, recommendation], ask
            )

    def answer_question(self, session_id: str, question_id: str | None) -> None:
        with self._transaction() as connection:
            session = self._require(connection, session_id)
            if question_id and session.get("pending_question") == question_id:
                session["pending_question"] = None
                self._save_session(connection, session)

    def save_case(
        self,
        session_id: str,
        case: CaseDraft,
        expected_version: int,
        request_id: str,
        *,
        source_run_id: str | None = None,
        generated: bool = False,
        generation: dict | None = None,
    ) -> dict:
        value = {
            "case": case.model_dump(),
            "expected_version": expected_version,
            "source_run_id": source_run_id,
            "generated": generated,
        }
        with self._transaction() as connection:
            session = self._require(connection, session_id)

            def save():
                if generation is not None:
                    current_generation = self._get(
                        connection, "exploration_generation", generation["id"]
                    )
                    if (
                        not current_generation
                        or current_generation.get("owner") != generation.get("owner")
                        or current_generation.get("status") != "running"
                    ):
                        raise ExplorationConflict("此提炼请求已由另一进程接管，请读取最新结果")
                if session["case_version"] != expected_version:
                    raise ExplorationConflict("用例已有更新，请加载最新版本后再保存")
                origin = (
                    self._get(connection, "exploration_run", source_run_id)
                    if source_run_id
                    else None
                )
                if source_run_id and (not origin or origin["exploration_id"] != session_id):
                    raise ValueError("用例依据的执行记录不属于当前探索")
                if generated and (
                    not origin or not (origin.get("outcome") or origin.get("verified_outcome"))
                ):
                    raise ValueError("自动提炼必须关联已保存实际结果的探索运行")
                version = expected_version + 1
                record = {
                    "exploration_id": session_id,
                    "version": version,
                    **case.model_dump(),
                    "source_run_id": source_run_id,
                    "generated": generated,
                    "created_at": time.time(),
                }
                self._put(connection, "exploration_case", f"{session_id}_{version:010d}", record)
                session["case_version"] = version
                self._save_session(connection, session)
                if generation is not None:
                    self._put(
                        connection,
                        "exploration_generation",
                        generation["id"],
                        {**generation, "status": "completed", "case": record, "lease_until": 0},
                    )
                return record

            return self._operation(connection, f"{session_id}:case", request_id, value, save)

    def case(self, session_id: str, version: int | None = None) -> dict | None:
        session = self.get(session_id)
        version = version if version is not None else session["case_version"]
        if version < 1:
            return None
        return self.database.get("exploration_case", f"{session_id}_{version:010d}")


explorations = ExplorationStore()