vent_ai_platform/app/db/repositories/document.py

from __future__ import annotations

import json
import html
import re
from datetime import datetime, timezone
from math import ceil

from sqlalchemy import delete, exists, func, or_, select
from sqlalchemy.orm import Session

from app.db.models.document import (
    BusinessDocument,
    DocumentExtraction,
    DocumentExtractionTemplate,
    DocumentProcessingRun,
)
from app.db.repositories.base import (
    create_model,
    delete_model,
    paginate,
    update_model,
)

ACTIVE_PROCESSING_STATUSES = ("pending", "running", "parsing", "extracting")
EMPTY_VALUE_MARKERS = {
    "null",
    "none",
    "nan",
    "--",
    "-",
    "未提供",
    "未填写",
    "待补充",
    "待填写",
    "暂无",
    "暂无数据",
}


def _is_meaningful_extraction_value(value: object) -> bool:
    text = str(value or "").strip()
    if not text:
        return False
    try:
        parsed = json.loads(text)
    except (TypeError, ValueError, json.JSONDecodeError):
        parsed = None
    if isinstance(parsed, (list, dict)) and not parsed:
        return False
    return text.lower() not in EMPTY_VALUE_MARKERS


def _template_field_keys(template: DocumentExtractionTemplate) -> set[str]:
    schema = template.extraction_schema or {}
    keys: set[str] = set()
    for group in schema.get("groups", []):
        if group.get("analysis_only"):
            continue
        for field in group.get("fields", []):
            key = str(field.get("key") or "").strip()
            if key:
                keys.add(key)
    return keys


def _month_start(value: datetime) -> datetime:
    return value.replace(day=1, hour=0, minute=0, second=0, microsecond=0)


def _previous_month_start(value: datetime) -> datetime:
    if value.month == 1:
        return value.replace(year=value.year - 1, month=12)
    return value.replace(month=value.month - 1)


def list_documents(
    db: Session,
    *,
    mine_id: str,
    category: str | None = None,
    keyword: str | None = None,
    parse_status: str | None = None,
    page: int = 1,
    page_size: int = 20,
) -> tuple[list[BusinessDocument], int]:
    statement = select(BusinessDocument).where(BusinessDocument.mine_id == mine_id)
    if category and category != "all":
        statement = statement.where(BusinessDocument.category == category)
    if parse_status:
        statement = statement.where(BusinessDocument.parse_status == parse_status)
    if keyword:
        query = f"%{keyword.strip()}%"
        statement = statement.where(
            or_(
                BusinessDocument.title.like(query),
                BusinessDocument.code.like(query),
                BusinessDocument.owner_name.like(query),
                BusinessDocument.summary.like(query),
            )
        )
    statement = statement.order_by(BusinessDocument.updated_at.desc())
    return paginate(db, statement, page, page_size)


def list_documents_for_export(
    db: Session,
    *,
    mine_id: str,
    category: str | None = None,
    keyword: str | None = None,
    parse_status: str | None = None,
) -> list[BusinessDocument]:
    statement = select(BusinessDocument).where(BusinessDocument.mine_id == mine_id)
    if category and category != "all":
        statement = statement.where(BusinessDocument.category == category)
    if parse_status:
        statement = statement.where(BusinessDocument.parse_status == parse_status)
    if keyword:
        query = f"%{keyword.strip()}%"
        statement = statement.where(
            or_(
                BusinessDocument.title.like(query),
                BusinessDocument.code.like(query),
                BusinessDocument.owner_name.like(query),
                BusinessDocument.summary.like(query),
            )
        )
    return list(
        db.scalars(
            statement.order_by(BusinessDocument.updated_at.desc()).limit(20000)
        ).all()
    )


def extraction_counts(
    db: Session,
    document_ids: list[str],
) -> dict[str, int]:
    if not document_ids:
        return {}
    rows = db.execute(
        select(
            DocumentExtraction.document_id,
            func.count(DocumentExtraction.id),
        )
        .where(DocumentExtraction.document_id.in_(document_ids))
        .group_by(DocumentExtraction.document_id)
    ).all()
    return {str(document_id): int(count) for document_id, count in rows}


def _clean_search_text(value: str) -> str:
    text = html.unescape(str(value or ""))
    text = re.sub(r"<[^>]+>", " ", text)
    text = re.sub(r"[ \t\r\f\v]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def _search_snippet(text: str, keyword: str, radius: int = 80) -> str:
    cleaned = _clean_search_text(text)
    if not cleaned:
        return ""
    position = cleaned.lower().find(keyword.lower())
    if position < 0:
        return cleaned[: radius * 2].strip()
    start = max(0, position - radius)
    end = min(len(cleaned), position + len(keyword) + radius)
    return f"{'…' if start > 0 else ''}{cleaned[start:end].strip()}{'…' if end < len(cleaned) else ''}"


def search_documents(
    db: Session,
    *,
    mine_id: str,
    query: str,
    category: str | None = None,
    page: int = 1,
    page_size: int = 10,
) -> tuple[list[dict], int]:
    keyword = query.strip()
    pattern = f"%{keyword}%"
    metadata_condition = or_(
        BusinessDocument.title.like(pattern),
        BusinessDocument.code.like(pattern),
        BusinessDocument.owner_name.like(pattern),
        BusinessDocument.department.like(pattern),
        BusinessDocument.file_name.like(pattern),
        BusinessDocument.summary.like(pattern),
    )
    parsed_text_condition = exists(
        select(DocumentProcessingRun.id).where(
            DocumentProcessingRun.document_id == BusinessDocument.id,
            DocumentProcessingRun.parsed_markdown.like(pattern),
        )
    )
    extraction_condition = exists(
        select(DocumentExtraction.id).where(
            DocumentExtraction.document_id == BusinessDocument.id,
            or_(
                DocumentExtraction.field_name.like(pattern),
                DocumentExtraction.field_value.like(pattern),
                DocumentExtraction.normalized_value.like(pattern),
                DocumentExtraction.evidence_text.like(pattern),
            ),
        )
    )
    statement = select(BusinessDocument).where(
        BusinessDocument.mine_id == mine_id,
        or_(
            metadata_condition,
            parsed_text_condition,
            extraction_condition,
        ),
    )
    if category and category != "all":
        statement = statement.where(BusinessDocument.category == category)
    statement = statement.order_by(BusinessDocument.updated_at.desc())
    documents, total = paginate(db, statement, page, page_size)

    hits: list[dict] = []
    for document in documents:
        score = 0.0
        snippets: list[dict[str, str]] = []
        metadata_values = [
            ("title", "资料名称", document.title),
            ("code", "文档编号", document.code or ""),
            ("owner", "资料责任人", document.owner_name or ""),
            ("department", "所属部门", document.department or ""),
            ("file_name", "文件名", document.file_name or ""),
            ("summary", "资料摘要", document.summary or ""),
        ]
        for source, label, value in metadata_values:
            if keyword.lower() in str(value or "").lower():
                score += {
                    "title": 100,
                    "code": 80,
                    "owner": 55,
                    "department": 45,
                    "file_name": 60,
                    "summary": 35,
                }[source]
                snippets.append(
                    {
                        "source": source,
                        "label": label,
                        "text": _search_snippet(value, keyword),
                    }
                )

        latest_run = get_latest_processing_run(db, document.id)
        if (
            latest_run
            and latest_run.parsed_markdown
            and keyword.lower() in latest_run.parsed_markdown.lower()
        ):
            occurrences = latest_run.parsed_markdown.lower().count(keyword.lower())
            score += 30 + min(occurrences, 10) * 3
            snippets.append(
                {
                    "source": "parsed_text",
                    "label": "解析全文",
                    "text": _search_snippet(latest_run.parsed_markdown, keyword),
                }
            )

        extraction_matches = list(
            db.scalars(
                select(DocumentExtraction)
                .where(
                    DocumentExtraction.document_id == document.id,
                    or_(
                        DocumentExtraction.field_name.like(pattern),
                        DocumentExtraction.field_value.like(pattern),
                        DocumentExtraction.normalized_value.like(pattern),
                        DocumentExtraction.evidence_text.like(pattern),
                    ),
                )
                .limit(5)
            ).all()
        )
        for extraction in extraction_matches:
            score += 20
            value = extraction.normalized_value or extraction.field_value or ""
            snippets.append(
                {
                    "source": "extraction",
                    "label": f"关键字段 · {extraction.field_name}",
                    "text": _search_snippet(f"{extraction.field_name}：{value}", keyword),
                }
            )
        hits.append(
            {
                "document": document,
                "score": round(score, 2),
                "snippets": snippets[:5],
            }
        )
    hits.sort(key=lambda item: item["score"], reverse=True)
    return hits, total


def get_document(db: Session, document_id: str) -> BusinessDocument | None:
    return db.get(BusinessDocument, document_id)


def list_extractions(db: Session, document_id: str) -> list[DocumentExtraction]:
    return list(
        db.scalars(
            select(DocumentExtraction)
            .where(DocumentExtraction.document_id == document_id)
            .order_by(DocumentExtraction.page_no, DocumentExtraction.field_name)
        ).all()
    )


def category_summary(db: Session, mine_id: str) -> list[tuple[str, int]]:
    rows = db.execute(
        select(BusinessDocument.category, func.count(BusinessDocument.id))
        .where(BusinessDocument.mine_id == mine_id)
        .group_by(BusinessDocument.category)
        .order_by(BusinessDocument.category)
    ).all()
    return [(str(category), int(count)) for category, count in rows]


def document_statistics(db: Session, mine_id: str) -> dict:
    now = datetime.now(timezone.utc).replace(tzinfo=None)
    month_start = _month_start(now)
    previous_month_start = _previous_month_start(month_start)

    total_documents = int(
        db.scalar(
            select(func.count(BusinessDocument.id)).where(
                BusinessDocument.mine_id == mine_id
            )
        )
        or 0
    )
    documents_added_this_month = int(
        db.scalar(
            select(func.count(BusinessDocument.id)).where(
                BusinessDocument.mine_id == mine_id,
                BusinessDocument.created_at >= month_start,
            )
        )
        or 0
    )
    documents = list(
        db.scalars(
            select(BusinessDocument).where(BusinessDocument.mine_id == mine_id)
        ).all()
    )
    document_ids = [document.id for document in documents]
    extraction_rows = db.execute(
        select(
            DocumentExtraction.document_id,
            DocumentExtraction.field_key,
            DocumentExtraction.created_at,
            DocumentExtraction.verified,
            DocumentExtraction.field_value,
            DocumentExtraction.edited,
            DocumentExtraction.extraction_correct,
            DocumentExtraction.edit_judgement_status,
        )
        .join(
            BusinessDocument,
            DocumentExtraction.document_id == BusinessDocument.id,
        )
        .where(BusinessDocument.mine_id == mine_id)
    ).all()

    extracted_fields = len(extraction_rows)
    latest_template_by_document: dict[str, str | None] = {}
    if document_ids:
        processing_runs = db.execute(
            select(
                DocumentProcessingRun.document_id,
                DocumentProcessingRun.template_id,
            )
            .where(DocumentProcessingRun.document_id.in_(document_ids))
            .order_by(DocumentProcessingRun.created_at.desc())
        ).all()
        for run in processing_runs:
            latest_template_by_document.setdefault(run.document_id, run.template_id)
    templates = list(db.scalars(select(DocumentExtractionTemplate)).all())
    template_by_id = {template.id: template for template in templates}
    template_by_code = {template.code: template for template in templates}
    template_by_type = {template.document_type: template for template in templates}
    rows_by_document: dict[str, dict[str, object]] = {}
    for row in extraction_rows:
        rows_by_document.setdefault(row.document_id, {})[row.field_key] = row

    expected_fields = 0
    valid_fields = 0
    invalid_fields = 0
    missing_fields = 0
    for document in documents:
        actual_rows = rows_by_document.get(document.id, {})
        template_id = latest_template_by_document.get(document.id)
        template = (
            template_by_id.get(template_id)
            if template_id
            else template_by_code.get(document.document_type)
            or template_by_type.get(document.document_type)
        )
        expected_keys = (
            _template_field_keys(template)
            if template is not None
            else set(actual_rows)
        )
        if not expected_keys:
            expected_keys = set(actual_rows)
        expected_fields += len(expected_keys)
        for key in expected_keys:
            row = actual_rows.get(key)
            if row is None:
                missing_fields += 1
            elif _is_meaningful_extraction_value(row.field_value):
                valid_fields += 1
            else:
                invalid_fields += 1

    verified_fields = sum(1 for row in extraction_rows if bool(row.verified))
    edited_fields = sum(1 for row in extraction_rows if bool(row.edited))
    pending_judgements = sum(
        1 for row in extraction_rows if row.edit_judgement_status == "pending"
    )
    incorrect_fields = sum(
        1
        for row in extraction_rows
        if row.edit_judgement_status != "pending"
        and not bool(row.extraction_correct)
    )
    accuracy_fields = extracted_fields - pending_judgements
    extraction_accuracy = (
        (accuracy_fields - incorrect_fields) / accuracy_fields * 100
        if accuracy_fields
        else 0.0
    )
    current_month_rows = [
        row
        for row in extraction_rows
        if row.created_at
        and row.created_at >= month_start
        and row.edit_judgement_status != "pending"
    ]
    previous_month_rows = [
        row
        for row in extraction_rows
        if row.created_at
        and previous_month_start <= row.created_at < month_start
        and row.edit_judgement_status != "pending"
    ]
    current_month_accuracy = (
        sum(1 for row in current_month_rows if row.extraction_correct)
        / len(current_month_rows)
        * 100
        if current_month_rows
        else 0.0
    )
    previous_month_accuracy = (
        sum(1 for row in previous_month_rows if row.extraction_correct)
        / len(previous_month_rows)
        * 100
        if previous_month_rows
        else 0.0
    )

    active_documents = list(
        db.scalars(
            select(BusinessDocument).where(
                BusinessDocument.mine_id == mine_id,
                BusinessDocument.parse_status.in_(ACTIVE_PROCESSING_STATUSES),
            )
        ).all()
    )
    estimated_seconds = 0.0
    for document in active_documents:
        if document.row_count:
            total_seconds = max(180.0, float(document.row_count) * 0.375)
        elif document.page_count:
            total_seconds = max(180.0, float(document.page_count) * 54)
        else:
            total_seconds = max(
                180.0,
                (float(document.file_size or 0) / (1024 * 1024)) * 48,
            )
        progress = max(0, min(int(document.parse_progress or 0), 100))
        estimated_seconds += total_seconds * (1 - progress / 100)

    field_completeness = (
        valid_fields / expected_fields * 100 if expected_fields else 0.0
    )
    accuracy_change = (
        current_month_accuracy - previous_month_accuracy
        if current_month_rows and previous_month_rows
        else 0.0
    )
    return {
        "total_documents": total_documents,
        "documents_added_this_month": documents_added_this_month,
        "extracted_fields": extracted_fields,
        "field_completeness": round(field_completeness, 1),
        "extraction_accuracy": round(extraction_accuracy, 1),
        "accuracy_change": round(accuracy_change, 1),
        "pending_tasks": len(active_documents),
        "estimated_minutes": ceil(estimated_seconds / 60) if estimated_seconds else 0,
        "verified_fields": verified_fields,
        "edited_fields": edited_fields,
        "incorrect_fields": incorrect_fields,
        "expected_fields": expected_fields,
        "valid_fields": valid_fields,
        "invalid_fields": invalid_fields,
        "missing_fields": missing_fields,
        "pending_judgements": pending_judgements,
        "accuracy_fields": accuracy_fields,
        "updated_at": now,
    }


def create_document(db: Session, data: dict) -> BusinessDocument:
    return create_model(db, BusinessDocument, data)


def update_document(
    db: Session,
    document: BusinessDocument,
    data: dict,
) -> BusinessDocument:
    return update_model(db, document, data)


def delete_document(db: Session, document: BusinessDocument) -> None:
    delete_model(db, document)


def create_extraction(db: Session, data: dict) -> DocumentExtraction:
    return create_model(db, DocumentExtraction, data)


def get_extraction(
    db: Session,
    extraction_id: str,
) -> DocumentExtraction | None:
    return db.get(DocumentExtraction, extraction_id)


def update_extraction(
    db: Session,
    extraction: DocumentExtraction,
    data: dict,
) -> DocumentExtraction:
    return update_model(db, extraction, data)


def delete_extractions(db: Session, document_id: str) -> None:
    db.execute(
        delete(DocumentExtraction).where(DocumentExtraction.document_id == document_id)
    )
    db.commit()


def list_templates(
    db: Session,
    document_type: str | None = None,
) -> list[DocumentExtractionTemplate]:
    statement = select(DocumentExtractionTemplate).where(
        DocumentExtractionTemplate.enabled.is_(True)
    )
    if document_type:
        statement = statement.where(
            DocumentExtractionTemplate.document_type == document_type
        )
    return list(
        db.scalars(
            statement.order_by(DocumentExtractionTemplate.document_type)
        ).all()
    )


def create_processing_run(
    db: Session,
    data: dict,
) -> DocumentProcessingRun:
    return create_model(db, DocumentProcessingRun, data)


def get_processing_run(
    db: Session,
    run_id: str,
) -> DocumentProcessingRun | None:
    return db.get(DocumentProcessingRun, run_id)


def get_latest_processing_run(
    db: Session,
    document_id: str,
) -> DocumentProcessingRun | None:
    return db.scalar(
        select(DocumentProcessingRun)
        .where(DocumentProcessingRun.document_id == document_id)
        .order_by(DocumentProcessingRun.created_at.desc())
        .limit(1)
    )


def _processing_duration_seconds(
    *,
    status: str,
    started_at: datetime | None,
    finished_at: datetime | None,
    created_at: datetime | None,
) -> float | None:
    def naive(value: datetime | None) -> datetime | None:
        if value is None:
            return None
        return value.replace(tzinfo=None) if value.tzinfo else value

    started = naive(started_at or created_at)
    if started is None:
        return None
    if status in {"completed", "failed"}:
        finished = naive(finished_at)
        if finished is None:
            return None
        return max(0.0, round((finished - started).total_seconds(), 3))
    if status in {"running", "parsing", "extracting"}:
        now = datetime.now(timezone.utc).replace(tzinfo=None)
        return max(0.0, round((now - started).total_seconds(), 3))
    return None


def list_processing_tasks(
    db: Session,
    *,
    mine_id: str,
    limit: int = 8,
) -> list[dict]:
    latest_created_at = (
        select(func.max(DocumentProcessingRun.created_at))
        .where(DocumentProcessingRun.document_id == BusinessDocument.id)
        .correlate(BusinessDocument)
        .scalar_subquery()
    )
    rows = db.execute(
        select(
            DocumentProcessingRun.id,
            DocumentProcessingRun.document_id,
            DocumentProcessingRun.status,
            DocumentProcessingRun.stage,
            DocumentProcessingRun.progress,
            DocumentProcessingRun.document_type.label("run_document_type"),
            DocumentProcessingRun.parser_engine,
            DocumentProcessingRun.error,
            DocumentProcessingRun.started_at,
            DocumentProcessingRun.finished_at,
            DocumentProcessingRun.created_at,
            DocumentProcessingRun.updated_at,
            BusinessDocument.title,
            BusinessDocument.file_name,
            BusinessDocument.file_ext,
            BusinessDocument.page_count,
            BusinessDocument.row_count,
            BusinessDocument.document_type.label("document_document_type"),
            BusinessDocument.parse_error,
        )
        .join(
            BusinessDocument,
            BusinessDocument.id == DocumentProcessingRun.document_id,
        )
        .where(
            BusinessDocument.mine_id == mine_id,
            DocumentProcessingRun.created_at == latest_created_at,
        )
        .order_by(DocumentProcessingRun.created_at.desc())
        .limit(limit)
    ).all()
    document_ids = [row.document_id for row in rows]
    counts: dict[str, int] = {}
    if document_ids:
        counts = {
            str(document_id): int(count)
            for document_id, count in db.execute(
                select(
                    DocumentExtraction.document_id,
                    func.count(DocumentExtraction.id),
                )
                .where(DocumentExtraction.document_id.in_(document_ids))
                .group_by(DocumentExtraction.document_id)
            ).all()
        }
    return [
        {
            "id": row.id,
            "document_id": row.document_id,
            "title": row.title,
            "file_name": row.file_name,
            "file_ext": row.file_ext,
            "status": row.status,
            "stage": row.stage,
            "progress": int(row.progress or 0),
            "document_type": row.run_document_type or row.document_document_type,
            "parser_engine": row.parser_engine,
            "page_count": row.page_count,
            "row_count": row.row_count,
            "field_count": counts.get(str(row.document_id), 0),
            "error": row.error or row.parse_error,
            "processing_seconds": _processing_duration_seconds(
                status=row.status,
                started_at=row.started_at,
                finished_at=row.finished_at,
                created_at=row.created_at,
            ),
            "started_at": row.started_at,
            "created_at": row.created_at,
            "updated_at": row.updated_at,
        }
        for row in rows
    ]