alert-contract-expiring-hr/scripts/db/repository.py

"""CRUD 仓储：contracts / performance_records / alerts / renewal_tasks / dispatch_log。

所有方法接收一个外部传入的 sqlite3.Connection，调用方负责事务边界。
"""

from __future__ import annotations

import json
import sqlite3
from datetime import date, datetime
from typing import Any, Iterable


# ---------- contracts ----------

def list_active_contracts(conn: sqlite3.Connection) -> list[sqlite3.Row]:
    """列出 status='可续聘' 的合同。"""
    return conn.execute(
        "SELECT * FROM contracts WHERE status = '可续聘' ORDER BY contract_end ASC"
    ).fetchall()


def get_contract(conn: sqlite3.Connection, contract_id: str) -> sqlite3.Row | None:
    return conn.execute(
        "SELECT * FROM contracts WHERE contract_id = ?", (contract_id,)
    ).fetchone()


def update_contract_status(
    conn: sqlite3.Connection,
    contract_id: str,
    status: str,
    material_id: str | None = None,
) -> int:
    """更新合同状态。material_id 仅在 status='已续签' 时写入。"""
    renewed_at = datetime.now().isoformat(timespec="seconds") if status == "已续签" else None
    cur = conn.execute(
        """
        UPDATE contracts
        SET status = ?,
            last_material_id = COALESCE(?, last_material_id),
            renewed_at = COALESCE(?, renewed_at)
        WHERE contract_id = ?
        """,
        (status, material_id, renewed_at, contract_id),
    )
    return cur.rowcount


# ---------- performance ----------

def list_performance_by_contract(
    conn: sqlite3.Connection, contract_id: str
) -> list[sqlite3.Row]:
    """某合同的业绩记录，按年份升序。"""
    return conn.execute(
        "SELECT * FROM performance_records WHERE contract_id = ? ORDER BY year ASC",
        (contract_id,),
    ).fetchall()


# 等次 → 分数映射（v1.1.1 sandbox card 解析）
_GRADE_TO_SCORE = {
    "优秀": 90.0,
    "称职": 75.0,
    "基本称职": 60.0,
    "不称职": 30.0,
    "不合格": 30.0,
}


def upsert_performance_record(
    conn: sqlite3.Connection,
    contract_id: str,
    year: int,
    grade: str,
) -> bool:
    """upsert 一条业绩记录（用 grade 映射分数）。

    sandbox card 文本不含真分数时使用。返回是否实际写入。
    """
    score = _GRADE_TO_SCORE.get(grade)
    if score is None:
        return False
    record_id = f"P-{contract_id}-{year}"
    conn.execute(
        """
        INSERT INTO performance_records
            (record_id, contract_id, year, performance_score,
             moral_teaching_score, research_score, teaching_score)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(record_id) DO UPDATE SET
            performance_score = excluded.performance_score
        """,
        (record_id, contract_id, year, score, score, score, score),
    )
    return True


def parse_sandbox_card_performance(
    conn: sqlite3.Connection,
    contract_id: str,
    card_text: str,
) -> int:
    """从 sandbox 卡片文本抽 (year, 等次) → upsert 到 performance_records。

    文本片段形如：
        year    result
        2024    称职
        2025    优秀
        2024年：等次「称职」。
    返回写入条数。
    """
    import re
    if not card_text:
        return 0

    written = 0
    # 模式 1：行格式 "2024\t称职"
    for m in re.finditer(r"^(\d{4})\s*[ \t]+(优秀|称职|基本称职|不称职|不合格)\s*$",
                         card_text, re.MULTILINE):
        year = int(m.group(1))
        grade = m.group(2)
        if upsert_performance_record(conn, contract_id, year, grade):
            written += 1

    # 模式 2：叙述式 "2024年：等次「称职」"
    for m in re.finditer(r"(\d{4})\s*年[：:]\s*等次[「\"']\s*(优秀|称职|基本称职|不称职|不合格)\s*[」\"']",
                         card_text):
        year = int(m.group(1))
        grade = m.group(2)
        if upsert_performance_record(conn, contract_id, year, grade):
            written += 1

    return written


# ---------- alerts ----------

def insert_alert(
    conn: sqlite3.Connection,
    alert_id: str,
    contract_id: str,
    level: str,
    days_to_expire: int,
    scan_id: str,
) -> None:
    conn.execute(
        """
        INSERT OR REPLACE INTO alerts
            (alert_id, contract_id, alert_level, days_to_expire, scan_id, generated_at, status)
        VALUES (?, ?, ?, ?, ?, ?, 'open')
        """,
        (
            alert_id,
            contract_id,
            level,
            days_to_expire,
            scan_id,
            datetime.now().isoformat(timespec="seconds"),
        ),
    )


def close_open_alerts_for_contract(conn: sqlite3.Connection, contract_id: str) -> int:
    """将某合同所有 open 状态的 alerts 关闭（关联 '已续签' 等终态）。"""
    cur = conn.execute(
        "UPDATE alerts SET status = 'closed' WHERE contract_id = ? AND status = 'open'",
        (contract_id,),
    )
    return cur.rowcount


def list_open_alerts_by_unit_month(
    conn: sqlite3.Connection, unit_id: str, month: str
) -> list[sqlite3.Row]:
    """按单位/月份列出关联合同的 open alerts（用于派发汇总）。"""
    sql = """
    SELECT a.alert_id, a.contract_id, a.alert_level, a.days_to_expire,
           c.contract_end, c.teacher_name, c.unit_name
    FROM alerts a
    JOIN contracts c ON c.contract_id = a.contract_id
    WHERE a.status = 'open'
      AND c.unit_id = ?
    """
    params: list[Any] = [unit_id]
    if month:
        sql += " AND substr(c.contract_end, 1, 7) = ?"
        params.append(month)
    sql += " ORDER BY c.contract_end ASC"
    return conn.execute(sql, params).fetchall()


# ---------- renewal_tasks ----------

def insert_task(
    conn: sqlite3.Connection,
    task_id: str,
    contract_id: str,
    alert_id: str,
    suggestion: str,
    reason: str,
    summary: str,
    checklist: list[str],
    decision_source: str = "rules",
    materials_summary: str = "",
) -> None:
    conn.execute(
        """
        INSERT OR REPLACE INTO renewal_tasks
            (task_id, contract_id, alert_id, ai_suggestion, suggestion_reason,
             performance_summary, materials_checklist, generated_at, status,
             decision_source, materials_summary)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)
        """,
        (
            task_id,
            contract_id,
            alert_id,
            suggestion,
            reason,
            summary,
            json.dumps(checklist, ensure_ascii=False),
            datetime.now().isoformat(timespec="seconds"),
            decision_source,
            materials_summary,
        ),
    )


def get_task(conn: sqlite3.Connection, task_id: str) -> sqlite3.Row | None:
    return conn.execute(
        """
        SELECT t.*, c.teacher_name, c.unit_name, c.position, c.contract_type,
               c.contract_start, c.contract_end, c.employee_id
        FROM renewal_tasks t
        JOIN contracts c ON c.contract_id = t.contract_id
        WHERE t.task_id = ?
        """,
        (task_id,),
    ).fetchone()


def get_task_full(conn: sqlite3.Connection, task_id: str) -> sqlite3.Row | None:
    """v1.1.1：拉取 task + 关联合同 + alert（含 days_to_expire/level）。"""
    return conn.execute(
        """
        SELECT t.*, c.teacher_name, c.unit_name, c.position, c.contract_type,
               c.contract_start, c.contract_end, c.employee_id, c.unit_id,
               a.days_to_expire AS alert_days, a.alert_level AS alert_level
        FROM renewal_tasks t
        JOIN contracts c ON c.contract_id = t.contract_id
        LEFT JOIN alerts a ON a.alert_id = t.alert_id
        WHERE t.task_id = ?
        """,
        (task_id,),
    ).fetchone()


# ---------- dispatch_log ----------

def insert_dispatch_log(
    conn: sqlite3.Connection,
    log_id: str,
    unit_id: str,
    month: str,
    task_ids: list[str],
    payload_path: str,
    status: str = "dry_run",
) -> None:
    conn.execute(
        """
        INSERT INTO dispatch_log
            (log_id, unit_id, month, task_ids, dispatched_at, payload_path, status)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        """,
        (
            log_id,
            unit_id,
            month,
            json.dumps(task_ids, ensure_ascii=False),
            datetime.now().isoformat(timespec="seconds"),
            payload_path,
            status,
        ),
    )


def list_dispatch_logs(conn: sqlite3.Connection) -> list[sqlite3.Row]:
    return conn.execute(
        "SELECT * FROM dispatch_log ORDER BY dispatched_at DESC"
    ).fetchall()


# ---------- materials_content（v1.1.1 新增） ----------

def upsert_contract_from_csv_row(
    conn: sqlite3.Connection, row: dict[str, str]
) -> None:
    """从 CSV 字典 upsert 一份合同。CSV 列名必须与 contracts 表列一致（或常见别名）。

    缺字段兜底值：
      employee_id      = contract_id
      teacher_name     = ""
      unit_id          = "U-UNKNOWN"
      unit_name        = ""
      position         = ""
      contract_type    = ""
      contract_start   = "1970-01-01"
      contract_end     = "2099-12-31"
      trial_period_end = contract_end
      status           = "可续聘"
    """
    cid = row.get("contract_id") or row.get("合同ID") or ""
    if not cid:
        raise ValueError("CSV row missing contract_id / 合同ID")
    conn.execute(
        """
        INSERT INTO contracts
            (contract_id, employee_id, teacher_name, unit_id, unit_name,
             position, contract_type, contract_start, contract_end,
             trial_period_end, status, last_material_id, renewed_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL)
        ON CONFLICT(contract_id) DO UPDATE SET
            teacher_name     = excluded.teacher_name,
            unit_name        = excluded.unit_name,
            position         = excluded.position,
            contract_type    = excluded.contract_type,
            contract_start   = excluded.contract_start,
            contract_end     = excluded.contract_end,
            trial_period_end = excluded.trial_period_end,
            status           = excluded.status
        """,
        (
            cid,
            row.get("employee_id") or row.get("工号") or cid,
            row.get("teacher_name") or row.get("教职工姓名") or "",
            row.get("unit_id") or row.get("单位ID") or "U-UNKNOWN",
            row.get("unit_name") or row.get("单位") or "",
            row.get("position") or row.get("岗位") or "",
            row.get("contract_type") or row.get("合同类型") or "",
            row.get("contract_start") or row.get("合同起始日") or "1970-01-01",
            row.get("contract_end") or row.get("合同到期日") or "2099-12-31",
            row.get("trial_period_end") or row.get("试用期截止") or row.get("contract_end") or "2099-12-31",
            row.get("status") or "可续聘",
        ),
    )


def list_materials_for_contract(
    conn: sqlite3.Connection, contract_id: str
) -> list[sqlite3.Row]:
    """列出某合同的全部材料文件。"""
    return conn.execute(
        "SELECT * FROM materials_content WHERE contract_id = ? ORDER BY parsed_at ASC",
        (contract_id,),
    ).fetchall()


def insert_material(
    conn: sqlite3.Connection,
    contract_id: str,
    file_name: str,
    file_path: str,
    extracted_text: str,
    parse_method: str,
) -> None:
    conn.execute(
        """
        INSERT OR REPLACE INTO materials_content
            (contract_id, file_name, file_path, extracted_text,
             parse_method, parsed_at)
        VALUES (?, ?, ?, ?, ?, ?)
        """,
        (
            contract_id,
            file_name,
            file_path,
            extracted_text,
            parse_method,
            datetime.now().isoformat(timespec="seconds"),
        ),
    )


def get_material(
    conn: sqlite3.Connection, contract_id: str, file_name: str
) -> sqlite3.Row | None:
    return conn.execute(
        "SELECT * FROM materials_content WHERE contract_id = ? AND file_name = ?",
        (contract_id, file_name),
    ).fetchone()


def list_pending_materials(conn: sqlite3.Connection) -> list[sqlite3.Row]:
    """列出待解析的材料（extracted_text IS NULL）。"""
    return conn.execute(
        "SELECT * FROM materials_content "
        "WHERE extracted_text IS NULL OR parse_method = 'pending' "
        "ORDER BY parsed_at ASC"
    ).fetchall()


def summarize_materials(
    conn: sqlite3.Connection, contract_id: str
) -> tuple[str, list[str]]:
    """汇总某合同所有已解析材料的文本 + 文件名列表。

    返回 (summary_text, file_names)。
    """
    rows = list_materials_for_contract(conn, contract_id)
    if not rows:
        return ("", [])

    file_names: list[str] = []
    parts: list[str] = []
    for r in rows:
        text = r["extracted_text"] or ""
        if not text:
            continue
        file_names.append(r["file_name"])
        parts.append(f"### {r['file_name']}\n{text}")

    summary = "\n\n".join(parts)
    # 截断 16k 字避免 LLM context 爆栈
    if len(summary) > 16000:
        summary = summary[:16000] + "\n...(材料过长已截断)"
    return (summary, file_names)


# ---------- download_tasks（v1.1.18 新增） ----------


def enqueue_download_task(
    conn: sqlite3.Connection,
    contract_id: str,
    kind: str = "platform_button",
) -> str:
    """入队一个下载任务，返回 task_id。

    已存在 (contract_id, kind, status ∈ {pending, running}) 时跳过（返回已存在的 task_id）。
    """
    import uuid

    existing = conn.execute(
        """SELECT task_id FROM download_tasks
           WHERE contract_id = ? AND kind = ? AND status IN ('pending', 'running')
           LIMIT 1""",
        (contract_id, kind),
    ).fetchone()
    if existing:
        return existing["task_id"]
    task_id = f"DLQ-{uuid.uuid4().hex[:12]}"
    conn.execute(
        """
        INSERT INTO download_tasks
            (task_id, contract_id, kind, status, attempts, enqueued_at)
        VALUES (?, ?, ?, 'pending', 0, ?)
        """,
        (task_id, contract_id, kind, datetime.now().isoformat(timespec="seconds")),
    )
    return task_id


def _read_max_attempts(conn: sqlite3.Connection) -> int:
    """读 DOWNLOAD_QUEUE_MAX_ATTEMPTS env（默认 3）。"""
    import os
    return int(os.environ.get("DOWNLOAD_QUEUE_MAX_ATTEMPTS", "3"))


def claim_pending_tasks(
    conn: sqlite3.Connection, limit: int, ids_filter: list[str] | None = None,
) -> list[sqlite3.Row]:
    """原子领取一批 pending 任务，标记为 running。

    - limit: 最多领取条数
    - ids_filter: 若给出，仅从这些 task_id 中领
    """
    max_attempts = _read_max_attempts(conn)
    if ids_filter:
        placeholders = ",".join("?" for _ in ids_filter)
        rows = conn.execute(
            f"""SELECT * FROM download_tasks
                WHERE status = 'pending' AND attempts < ?
                AND task_id IN ({placeholders})
                ORDER BY enqueued_at ASC LIMIT ?""",
            [max_attempts, *ids_filter, limit],
        ).fetchall()
    else:
        rows = conn.execute(
            """SELECT * FROM download_tasks
               WHERE status = 'pending' AND attempts < ?
               ORDER BY enqueued_at ASC LIMIT ?""",
            (max_attempts, limit),
        ).fetchall()
    if not rows:
        return []
    now = datetime.now().isoformat(timespec="seconds")
    for r in rows:
        conn.execute(
            "UPDATE download_tasks SET status = 'running', started_at = ? WHERE task_id = ?",
            (now, r["task_id"]),
        )
    conn.commit()
    return rows


def mark_task_success(
    conn: sqlite3.Connection, task_id: str, result_path: str,
) -> None:
    conn.execute(
        """UPDATE download_tasks
           SET status = 'success', result_path = ?, finished_at = ?
           WHERE task_id = ?""",
        (result_path, datetime.now().isoformat(timespec="seconds"), task_id),
    )


def mark_task_failed(
    conn: sqlite3.Connection, task_id: str, error_msg: str,
) -> None:
    """失败一次：attempts+1，未超 max_attempts 则回退 pending，否则标 failed。"""
    max_attempts = _read_max_attempts(conn)
    now = datetime.now().isoformat(timespec="seconds")
    conn.execute(
        """UPDATE download_tasks
           SET status = CASE WHEN attempts + 1 >= ? THEN 'failed' ELSE 'pending' END,
               attempts = attempts + 1,
               error_msg = ?,
               finished_at = ?
           WHERE task_id = ?""",
        (max_attempts, error_msg, now, task_id),
    )


def list_recent_tasks(
    conn: sqlite3.Connection, limit: int = 50
) -> list[sqlite3.Row]:
    return conn.execute(
        "SELECT * FROM download_tasks ORDER BY enqueued_at DESC LIMIT ?",
        (limit,),
    ).fetchall()


def count_tasks_by_status(conn: sqlite3.Connection) -> dict[str, int]:
    rows = conn.execute(
        "SELECT status, COUNT(*) AS n FROM download_tasks GROUP BY status"
    ).fetchall()
    return {r["status"]: r["n"] for r in rows}