Agent/examples/content_finder/db/schedule.py

"""
定时任务相关数据库操作

demand_content: 原始检索内容库（name→query，suggestion→补充信息）
demand_find_task: 执行记录表，通过 demand_content_id 关联
"""

import logging
from datetime import datetime
from decimal import Decimal
from typing import Any, Dict, List, Optional

from .connection import get_connection

logger = logging.getLogger(__name__)

# 状态常量（与 demand_find_task 表一致）
STATUS_PENDING = 0
STATUS_RUNNING = 1
STATUS_SUCCESS = 2
STATUS_FAILED = 3


def get_next_unprocessed_demand() -> Optional[Dict[str, Any]]:
    """
    联表查询 demand_content 和 demand_find_task，找到创建最早且未处理的 demand_content。

    未处理定义：该 demand_content_id 在 demand_find_task 中尚无任何记录。
    已有任务则视为已跑过（含失败），不再被定时任务选中。

    Returns:
        {"demand_content_id": int, "query": str, "suggestion": Optional[str]} 或 None
    """
    sql = """
    SELECT dc.id AS demand_content_id,
           dc.name AS query,
           dc.suggestion AS suggestion
    FROM demand_content dc
    WHERE NOT EXISTS (
        SELECT 1 FROM demand_find_task t
        WHERE t.demand_content_id = dc.id
    )
    ORDER BY dc.id ASC
    LIMIT 1
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql)
            row = cur.fetchone()
            return dict(row) if row else None
    except Exception as e:
        logger.error(f"get_next_unprocessed_demand 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def get_first_running_task() -> Optional[Dict[str, Any]]:
    """
    查找 demand_find_task 中 status=STATUS_RUNNING(1) 的任务（理论上仅一条）。

    用于服务重启后恢复执行中的任务：联表取出 query（demand_content.name）、suggestion。

    Returns:
        {"demand_content_id": int, "query": str, "suggestion": Optional[str], "trace_id": str} 或 None
    """
    sql = """
    SELECT t.demand_content_id,
           t.trace_id,
           dc.name AS query,
           dc.suggestion AS suggestion
    FROM demand_find_task t
    INNER JOIN demand_content dc ON dc.id = t.demand_content_id
    WHERE t.status = %s
    ORDER BY t.id ASC
    LIMIT 1
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql, (STATUS_RUNNING,))
            row = cur.fetchone()
            return dict(row) if row else None
    except Exception as e:
        logger.error(f"get_first_running_task 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def get_latest_demand_task_oprate_is_open() -> Optional[int]:
    """
    读取 demand_task_oprate 中按 update_time 最新的一行的 is_open。

    - 1：允许定时寻找任务继续派发
    - 0：关闭定时派发
    - 无记录时返回 None，调用方可视为默认打开（与未建表/无数据时的兼容行为）

    Returns:
        is_open 整数值，或表为空时 None
    """
    sql = """
    SELECT is_open
    FROM demand_task_oprate
    ORDER BY update_time DESC, id DESC
    LIMIT 1
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql)
            row = cur.fetchone()
            if not row:
                return None
            val = row.get("is_open")
            return int(val) if val is not None else None
    except Exception as e:
        logger.error(f"get_latest_demand_task_oprate_is_open 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def get_latest_day_limit_coast() -> Optional[Decimal]:
    """
    读取 demand_task_oprate 中按 update_time 最新的一行的 day_limit_coast。

    Returns:
        Decimal(两位小数) 或 None（无记录或字段为空）
    """
    sql = """
    SELECT day_limit_coast
    FROM demand_task_oprate
    ORDER BY update_time DESC, id DESC
    LIMIT 1
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql)
            row = cur.fetchone()
            if not row:
                return None
            val = row.get("day_limit_coast")
            if val is None:
                return None
            return Decimal(str(val))
    except Exception as e:
        logger.error(f"get_latest_day_limit_coast 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def get_total_token_coast_between(start_time: datetime, end_time: datetime) -> Decimal:
    """
    统计 demand_find_task 在指定时间区间内的 token_coast 总和（基于 created_at 字段）。

    start_time <= created_at < end_time
    """
    sql = """
    SELECT COALESCE(SUM(token_coast), 0) AS total_coast
    FROM demand_find_task
    WHERE created_at >= %s
      AND created_at < %s
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql, (start_time, end_time))
            row = cur.fetchone() or {}
            val = row.get("total_coast") if row is not None else None
            if val is None:
                return Decimal("0")
            return Decimal(str(val))
    except Exception as e:
        logger.error(f"get_total_token_coast_between 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def get_one_today_unprocessed_demand(*, dt: int) -> Optional[Dict[str, Any]]:
    """
    从 demand_content 中取「当天 dt」且尚未在 demand_find_task 中出现过的 1 条需求，
    采用“分品类分层轮转”策略：

    - 先取每个品类分数最高的 1 条（各品类 top1），按 score 降序优先执行
    - 当各品类 top1 跑完后，再取各品类 top2，仍按 score 降序优先执行
    - 依此类推
    - dt 与表字段一致：一般为 YYYYMMDD 整数（如 20260402）

    Returns:
        {
            "demand_content_id": int,
            "query": str,
            "suggestion": Optional[str],
            "score": Any,
            "merge_leve2": Optional[str],
            "category_rank": int,
        } 或 None
    """
    sql = """
    SELECT x.demand_content_id,
           x.query,
           x.suggestion,
           x.score,
           x.merge_leve2,
           x.rn AS category_rank
    FROM (
        SELECT dc.id AS demand_content_id,
               dc.name AS query,
               dc.suggestion AS suggestion,
               dc.score AS score,
               dc.merge_leve2 AS merge_leve2,
               ROW_NUMBER() OVER (
                   PARTITION BY COALESCE(dc.merge_leve2, '')
                   ORDER BY dc.score DESC, dc.id DESC
               ) AS rn
        FROM demand_content dc
        WHERE dc.dt = %s
          AND NOT EXISTS (
            SELECT 1 FROM demand_find_task t
            WHERE t.demand_content_id = dc.id
          )
    ) x
    ORDER BY x.rn ASC, x.score DESC, x.demand_content_id DESC
    LIMIT 1
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql, (int(dt),))
            row = cur.fetchone()
            return dict(row) if row else None
    except Exception as e:
        logger.error(f"get_one_today_unprocessed_demand 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def get_daily_unprocessed_pool(
    *,
    total_limit: int = 20,
    per_category_limit: int = 5,
) -> List[Dict[str, Any]]:
    """
    生成“当天任务池”：
    - 先按 demand_content.merge_leve2 品类去重分组
    - 每个品类取 score 最高的 N 条（per_category_limit）
    - 全局最多取 M 条（total_limit）
    - 过滤已处理：demand_find_task 中存在任意记录则视为已跑过（含失败）
    """
    sql = """
    SELECT x.demand_content_id, x.query, x.suggestion, x.merge_leve2, x.score
    FROM (
        SELECT
            dc.id AS demand_content_id,
            dc.name AS query,
            dc.suggestion AS suggestion,
            dc.merge_leve2 AS merge_leve2,
            dc.score AS score,
            ROW_NUMBER() OVER (
                PARTITION BY COALESCE(dc.merge_leve2, '')
                ORDER BY dc.score DESC, dc.id DESC
            ) AS rn
        FROM demand_content dc
        WHERE NOT EXISTS (
            SELECT 1 FROM demand_find_task t
            WHERE t.demand_content_id = dc.id
        )
    ) x
    WHERE x.rn <= %s
    ORDER BY x.score DESC, x.demand_content_id DESC
    LIMIT %s
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql, (per_category_limit, total_limit))
            rows = cur.fetchall() or []
            return [dict(r) for r in rows]
    except Exception as e:
        logger.error(f"get_daily_unprocessed_pool 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def create_task_record(demand_content_id: int, trace_id: str = "", status: int = STATUS_PENDING) -> None:
    """
    在 demand_find_task 中新增一条记录。

    初始创建时 trace_id 可置为空字符串，任务完成后通过 update_task_on_complete 更新。

    若 uk_trace_demand 冲突：仅当原行 status 为失败时，将 trace_id/status 重置为本次插入值
    （用于上次失败且 trace_id 仍为空时，定时任务会再次选中同一需求，不能重复 INSERT 空 trace）。
    """
    sql = """
    INSERT INTO demand_find_task (trace_id, demand_content_id, status)
    VALUES (%s, %s, %s)
    ON DUPLICATE KEY UPDATE
        trace_id = IF(status = %s, VALUES(trace_id), trace_id),
        status = IF(status = %s, VALUES(status), status)
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(
                sql,
                (trace_id, demand_content_id, status, STATUS_FAILED, STATUS_FAILED),
            )
        logger.info(f"创建任务记录: demand_content_id={demand_content_id}")
    except Exception as e:
        logger.error(f"create_task_record 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def update_task_on_complete(
    demand_content_id: int,
    trace_id: str,
    status: int,
    token_coast: Optional[Decimal] = None,
) -> None:
    """
    任务完成后更新 trace_id 和 status。
    匹配 trace_id 为空字符串的记录（初始创建时的占位）。
    """
    sql = """
    UPDATE demand_find_task
    SET trace_id = %s,
        status = %s
    """
    params: list[Any] = [trace_id, status]

    if token_coast is not None:
        sql += ", token_coast = %s\n"
        params.append(token_coast)

    sql += "WHERE demand_content_id = %s AND trace_id = ''"
    params.append(demand_content_id)
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql, tuple(params))
        logger.info(
            "更新任务完成: demand_content_id=%s, trace_id=%s, status=%s, token_coast=%s",
            demand_content_id,
            trace_id,
            status,
            token_coast,
        )
    except Exception as e:
        logger.error("update_task_on_complete 失败: %s", e, exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def fetch_trace_ids_created_after(cutoff: datetime) -> list[str]:
    """
    查询 demand_find_task 中 created_at 晚于 cutoff 的去重 trace_id（排除空串）。

    cutoff 与表字段比较规则与 MySQL DATETIME/TIMESTAMP 一致。
    """
    sql = """
    SELECT DISTINCT trace_id
    FROM demand_find_task
    WHERE created_at > %s
      AND trace_id IS NOT NULL
      AND trace_id <> ''
    ORDER BY trace_id
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql, (cutoff,))
            rows = cur.fetchall() or []
            return [str(r["trace_id"]) for r in rows]
    except Exception as e:
        logger.error(f"fetch_trace_ids_created_after 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()


def update_task_status(trace_id: str, demand_content_id: int, status: int) -> None:
    """
    更新 demand_find_task 中指定记录的状态。

    trace_id 可为空字符串（任务尚未返回时，通过 demand_content_id 定位记录）。
    """
    sql = """
    UPDATE demand_find_task
    SET status = %s
    WHERE trace_id = %s AND demand_content_id = %s
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cur:
            cur.execute(sql, (status, trace_id, demand_content_id))
        logger.info(f"更新任务状态: trace_id={trace_id}, demand_content_id={demand_content_id}, status={status}")
    except Exception as e:
        logger.error(f"update_task_status 失败: {e}", exc_info=True)
        raise
    finally:
        if conn:
            conn.close()