"""持久任务队列（交付报告 P0-3）。

用 SQLite 表实现可靠队列语义，替代纯内存线程调度：
- 租约 lease：worker 领取任务时写入 lease_until，超时未续约视为 worker 死亡，任务可被其他 worker 重新领取；
- 重试退避：失败后按 attempt 计算退避（复用 hermes_core.jittered_backoff），next_run_at 之前不可领取；
- 死信队列：超过最大重试次数进入 dead 状态，运营可查看和手动重放；
- 心跳：worker 定期刷新 lease_until；服务重启时孤儿任务自动回收。

与现有 _ACTIVE_RUNS 内存调度共存：队列表是事实源，内存线程只是执行器。
"""
from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone

from sqlmodel import Session, select, update

from .hermes_core import jittered_backoff
from .models import now

logger = logging.getLogger("yawen.queue")

LEASE_SECONDS = 300          # 租约时长：5 分钟无心跳即视为 worker 死亡
MAX_ATTEMPTS = 3             # 含首次执行在内最多尝试次数
QUEUE_TABLE = "taskqueue"


def _utcnow() -> datetime:
    """统一使用无时区本地时间（与 models.now() 一致，队列内部时间戳互相比较）。"""
    return datetime.now()


def _parse(value: str) -> datetime | None:
    try:
        parsed = datetime.fromisoformat(value)
    except (ValueError, TypeError):
        return None
    if parsed.tzinfo is not None:
        parsed = parsed.replace(tzinfo=None)
    return parsed


def enqueue(session: Session, queue_model, *, task_key: str, payload: dict,
            update_existing: bool = False) -> int:
    """入队（幂等：同 task_key 的活跃任务不重复入队）。返回队列行 id。

    update_existing 用于 parse/compose/redo 这类确定性 campaign 任务：如果任务
    仍在 pending，用户刚修改了参数，则刷新 payload；running 中的任务不能被改。
    """
    existing = session.exec(select(queue_model).where(
        queue_model.task_key == task_key,
        queue_model.status.in_(["pending", "leased", "running"]),
    )).first()
    if existing:
        if update_existing and existing.status == "pending":
            existing.payload = payload
            existing.error = ""
            existing.next_run_at = now()
            existing.updated_at = now()
            session.add(existing)
            session.commit()
        return existing.id
    row = queue_model(task_key=task_key, payload=payload, status="pending",
                      attempt=0, max_attempts=MAX_ATTEMPTS, next_run_at=now())
    session.add(row)
    session.commit()
    session.refresh(row)
    return row.id


def lease_next(session: Session, queue_model, worker_id: str,
               task_key_prefix: str = "") -> object | None:
    """领取一个到期任务，并以条件 UPDATE 原子认领。

    旧实现先读取候选行、再提交修改；两个 worker 在读取与提交之间会
    同时看到同一行，导致重复执行。这里仍保留轻量候选查询以维持优先级，
    但真正认领通过 ``UPDATE ... WHERE`` 校验状态和租约，失败时重试下一个
    候选。SQLite 的单条 UPDATE 在写入时持有数据库写锁，因此不会产生双重
    认领。
    """
    for _ in range(8):
        current = now()
        # 仅取少量候选，避免队列很大时把全部 payload 读进内存。时间戳使用
        # 固定的 YYYY-MM-DD HH:MM:SS 格式，SQLite 文本比较与 _parse 语义一致。
        query = select(queue_model).where(
            queue_model.status.in_(["pending", "running"]),
        )
        # Worker 可按 task_key 前缀领取自己的任务类型。DSH 内容生产执行器
        # 只处理 campaign-run:*，避免误领取训练导图等其他后台任务。
        if task_key_prefix:
            query = query.where(
                queue_model.task_key.like(task_key_prefix.rstrip("%") + "%")
            )
        rows = session.exec(query.order_by(queue_model.id).limit(16)).all()
        candidate = None
        for row in rows:
            if row.status == "pending":
                next_at = _parse(row.next_run_at)
                if next_at and next_at > _parse(current):
                    continue
            else:  # running：租约超时才能抢（worker 死亡的孤儿任务）
                lease_until = _parse(row.lease_until or "")
                if lease_until and lease_until > _parse(current):
                    continue
            candidate = row
            break
        if candidate is None:
            return None

        deadline = _lease_deadline()
        # 把候选的当前状态/租约再次放入 WHERE，防止另一个 worker 已经抢走。
        # 对 pending 同时检查 next_run_at，避免失败重试被提前领取。
        # Match the exact timestamp observed in the candidate row instead of
        # comparing strings in SQL. Tests and older rows may include fractional
        # seconds while ``now()`` uses second precision; Python's _parse above
        # handles both formats correctly.
        pending_due = (
            (queue_model.status == "pending")
            & (queue_model.next_run_at == candidate.next_run_at)
        )
        orphan_due = (
            (queue_model.status == "running")
            & (queue_model.lease_until == candidate.lease_until)
        )
        statement = update(queue_model).where(
            queue_model.id == candidate.id,
            pending_due | orphan_due,
        ).values(
            status="running",
            attempt=(queue_model.attempt + 1),
            worker_id=worker_id,
            lease_until=deadline,
            updated_at=current,
        )
        result = session.exec(statement)
        # SQLAlchemy's rowcount is available for SQLite UPDATE statements. A
        # zero count means another worker won the race; rollback and retry.
        if result.rowcount != 1:
            session.rollback()
            continue
        session.commit()
        claimed = session.get(queue_model, candidate.id)
        if claimed is not None:
            return claimed
    return None


def _lease_deadline() -> str:
    return (_utcnow() + timedelta(seconds=LEASE_SECONDS)).isoformat()


def heartbeat(session: Session, queue_model, task_id: int, worker_id: str) -> bool:
    """worker 心跳续租；任务已不在本 worker 名下时返回 False。"""
    row = session.get(queue_model, task_id)
    if not row or row.status != "running" or row.worker_id != worker_id:
        return False
    row.lease_until = _lease_deadline()
    row.updated_at = now()
    session.add(row)
    session.commit()
    return True


def complete(session: Session, queue_model, task_id: int, worker_id: str) -> bool:
    row = session.get(queue_model, task_id)
    if not row or row.status != "running" or row.worker_id != worker_id:
        return False
    row.status = "done"
    row.error = ""
    row.lease_until = ""
    row.updated_at = now()
    session.add(row)
    session.commit()
    return True


def fail(session: Session, queue_model, task_id: int, worker_id: str,
         error: str) -> str:
    """任务失败：可重试则退避后重新排队；超过次数进死信。返回最终状态。"""
    row = session.get(queue_model, task_id)
    if not row or row.worker_id != worker_id:
        return "unknown"
    row.error = str(error)[:500]
    if (row.attempt or 0) >= (row.max_attempts or MAX_ATTEMPTS):
        row.status = "dead"
        logger.error("任务 %s 进入死信队列：%s", row.task_key, row.error[:120])
    else:
        delay = min(24.0, jittered_backoff(row.attempt))
        retry_at = (_utcnow() + timedelta(seconds=delay)).isoformat()
        row.status = "pending"
        row.next_run_at = retry_at
        row.lease_until = ""
        logger.warning("任务 %s 第 %s 次失败，%.1f 秒后重试：%s",
                       row.task_key, row.attempt, delay, row.error[:120])
    row.updated_at = now()
    session.add(row)
    session.commit()
    return row.status


def requeue_dead(session: Session, queue_model, task_id: int) -> bool:
    """手动重放死信任务。"""
    row = session.get(queue_model, task_id)
    if not row or row.status != "dead":
        return False
    row.status = "pending"
    row.attempt = 0
    row.next_run_at = now()
    row.error = ""
    row.updated_at = now()
    session.add(row)
    session.commit()
    return True


def queue_stats(session: Session, queue_model) -> dict:
    rows = session.exec(select(queue_model)).all()
    by_status: dict[str, int] = {"pending": 0, "running": 0, "done": 0, "dead": 0}
    for row in rows:
        by_status[row.status] = by_status.get(row.status, 0) + 1
    return {"total": len(rows), "by_status": by_status,
            "done": by_status.get("done", 0),
            "dead": by_status.get("dead", 0),
            "pending": by_status.get("pending", 0),
            "running": by_status.get("running", 0)}


__all__ = ["enqueue", "lease_next", "claim_by_key", "heartbeat", "complete", "fail",
           "requeue_dead", "queue_stats", "LEASE_SECONDS", "MAX_ATTEMPTS"]


def claim_by_key(session: Session, queue_model, task_key: str, worker_id: str) -> object | None:
    """Claim a specific durable job for the inline API worker.

    A running row is only reused by the same owner or stolen after its lease
    expires; otherwise another live executor owns it.
    """
    row = session.exec(select(queue_model).where(queue_model.task_key == task_key).order_by(queue_model.id.desc())).first()
    if row is None:
        return None
    current = now()
    expired = False
    if row.status == "running":
        deadline = _parse(row.lease_until or "")
        expired = bool(deadline and deadline <= _parse(current))
    if row.status == "pending" or (row.status == "running" and expired):
        row.status = "running"
        row.attempt = (row.attempt or 0) + 1
        row.worker_id = worker_id
        row.lease_until = _lease_deadline()
        row.updated_at = current
        session.add(row)
        session.commit()
        session.refresh(row)
        return row
    if row.status == "running" and row.worker_id == worker_id:
        return row
    return None