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()


# ---------- 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 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)