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