arm-core/api/src/repositories/statements.py

"""Таблицы `financial_statement_*` (Срез 3b): вставка формы при маппинге, очистка файла, чтение сумм и флагов значений.
Только SQL, без правил: правила — `domain/statement_mapping.py`, `domain/statement_checks.py`."""

import datetime
import uuid
from collections.abc import Sequence
from decimal import Decimal
from typing import Any

from sqlalchemy import delete, insert, select
from sqlalchemy.ext.asyncio import AsyncSession

from src.domain.statement_mapping import StatementDraft
from src.models.enums import DocumentStatus, DocumentType, FileStatus
from src.models.files import Document, StoredFile
from src.models.statements import (
    FinancialStatementColumn,
    FinancialStatementLine,
    FinancialStatementMeta,
    FinancialStatementValue,
)

__all__ = [
    "BATCH",
    "delete_file_statements",
    "get_meta",
    "insert_statement",
    "load_form",
    "lock_cells",
    "pair_amounts",
    "statement_amounts",
    "value_flags",
]

BATCH = 1000  # значений на один INSERT: 8 колонок × 1000 < 32767 параметров asyncpg


async def insert_statement(
    session: AsyncSession, *, document_id: uuid.UUID, job_id: uuid.UUID, draft: StatementDraft
) -> FinancialStatementMeta:
    """Шапка, колонки, строки и значения формы из черновика (без commit); возвращает шапку."""
    meta = FinancialStatementMeta(
        id=uuid.uuid4(),
        document_id=document_id,
        recognition_job_id=job_id,
        statement_type=draft.statement_type,
        **draft.header,
    )
    session.add(meta)
    await session.flush()
    column_ids: dict[int, uuid.UUID] = {}
    columns: list[dict[str, Any]] = []
    for column_index, label, period_date, is_current in draft.columns:
        column_ids[column_index] = uuid.uuid4()
        columns.append(
            {
                "id": column_ids[column_index],
                "meta_id": meta.id,
                "column_index": column_index,
                "label": label,
                "period_date": period_date,
                "is_current": is_current,
            }
        )
    if columns:
        await session.execute(insert(FinancialStatementColumn).values(columns))
    line_ids: dict[str, uuid.UUID] = {}
    lines: list[dict[str, Any]] = []
    for line in draft.lines:
        line_ids[line.line_code] = uuid.uuid4()
        lines.append(
            {
                "id": line_ids[line.line_code],
                "meta_id": meta.id,
                "line_code": line.line_code,
                "line_code_original_value": line.line_code,
                "line_code_confidence": line.line_code_confidence,
                "line_code_modified": False,
                "line_name": line.line_name,
                "line_name_original_value": line.line_name,
                "line_name_confidence": line.line_name_confidence,
                "line_name_modified": False,
            }
        )
    for start in range(0, len(lines), BATCH):
        await session.execute(insert(FinancialStatementLine).values(lines[start : start + BATCH]))
    values = [
        {
            "id": uuid.uuid4(),
            "line_id": line_ids[value.line_code],
            "column_id": column_ids[value.column_index],
            "amount": value.amount,
            "string_value": value.string_value,
            "confidence": value.confidence,
            "is_recognized": value.is_recognized,
            "is_modified": False,
        }
        for value in draft.values
    ]
    for start in range(0, len(values), BATCH):
        await session.execute(insert(FinancialStatementValue).values(values[start : start + BATCH]))
    return meta


async def delete_file_statements(session: AsyncSession, file_id: uuid.UUID) -> None:
    """Значения, строки, колонки и шапки форм всех документов файла (в порядке FK)."""
    documents = select(Document.id).where(Document.stored_file_id == file_id).scalar_subquery()
    metas = select(FinancialStatementMeta.id).where(FinancialStatementMeta.document_id.in_(documents)).scalar_subquery()
    lines = select(FinancialStatementLine.id).where(FinancialStatementLine.meta_id.in_(metas)).scalar_subquery()
    await session.execute(delete(FinancialStatementValue).where(FinancialStatementValue.line_id.in_(lines)))
    await session.execute(delete(FinancialStatementLine).where(FinancialStatementLine.meta_id.in_(metas)))
    await session.execute(delete(FinancialStatementColumn).where(FinancialStatementColumn.meta_id.in_(metas)))
    await session.execute(delete(FinancialStatementMeta).where(FinancialStatementMeta.document_id.in_(documents)))


async def has_lines(session: AsyncSession, meta_id: uuid.UUID) -> bool:
    """В форме есть строки `financial_statement_line` (единое определение «есть строки» для статуса D-6)."""
    stmt = select(select(FinancialStatementLine.id).where(FinancialStatementLine.meta_id == meta_id).exists())
    return bool((await session.execute(stmt)).scalar_one())


async def statement_amounts(
    session: AsyncSession, meta_id: uuid.UUID
) -> tuple[list[int], list[tuple[str, dict[int, Decimal | None]]]]:
    """Индексы колонок и `(код строки, {колонка: сумма})` по возрастанию кода — вход `run_checks`. Порядок строк
    формы в БД не хранится (DBML не хранит), поэтому порядок строк — лексикографический порядок кодов."""
    columns = (
        await session.execute(
            select(FinancialStatementColumn.id, FinancialStatementColumn.column_index)
            .where(FinancialStatementColumn.meta_id == meta_id)
            .order_by(FinancialStatementColumn.column_index)
        )
    ).all()
    index_by_id = {column_id: index for column_id, index in columns}
    rows = (
        await session.execute(
            select(FinancialStatementLine.line_code, FinancialStatementValue.column_id, FinancialStatementValue.amount)
            .join(FinancialStatementValue, FinancialStatementValue.line_id == FinancialStatementLine.id)
            .where(FinancialStatementLine.meta_id == meta_id)
            .order_by(FinancialStatementLine.line_code)
        )
    ).all()
    by_code: dict[str, dict[int, Decimal | None]] = {}
    for code, column_id, amount in rows:
        by_code.setdefault(code or "", {})[index_by_id[column_id]] = amount
    return [index for _, index in columns], list(by_code.items())


async def value_flags(session: AsyncSession, meta_id: uuid.UUID) -> list[tuple[bool, bool | None, Decimal | None]]:
    """`(is_modified, is_recognized, confidence)` всех значений формы — вход `values_ok`."""
    rows = await session.execute(
        select(
            FinancialStatementValue.is_modified,
            FinancialStatementValue.is_recognized,
            FinancialStatementValue.confidence,
        )
        .join(FinancialStatementLine, FinancialStatementLine.id == FinancialStatementValue.line_id)
        .where(FinancialStatementLine.meta_id == meta_id)
    )
    return [(bool(modified), recognized, confidence) for modified, recognized, confidence in rows]


async def lock_cells(
    session: AsyncSession, meta_id: uuid.UUID, cells: Sequence[tuple[uuid.UUID, uuid.UUID]]
) -> dict[tuple[uuid.UUID, uuid.UUID], FinancialStatementValue]:
    """Значения формы `meta_id` по `(line_id, column_id)` под `FOR UPDATE` с перечитыванием; чужие ячейки (строка не
    этой формы или колонка не её) в словарь не попадают."""
    if not cells:
        return {}
    line_ids = {line_id for line_id, _ in cells}
    stmt = (
        select(FinancialStatementValue)
        .join(FinancialStatementLine, FinancialStatementLine.id == FinancialStatementValue.line_id)
        .join(FinancialStatementColumn, FinancialStatementColumn.id == FinancialStatementValue.column_id)
        .where(
            FinancialStatementLine.meta_id == meta_id,
            FinancialStatementColumn.meta_id == meta_id,
            FinancialStatementValue.line_id.in_(line_ids),
        )
        .order_by(FinancialStatementValue.id)
        .with_for_update(of=FinancialStatementValue)
        .execution_options(populate_existing=True)
    )
    wanted = set(cells)
    return {
        (value.line_id, value.column_id): value
        for value in (await session.execute(stmt)).scalars()
        if (value.line_id, value.column_id) in wanted
    }


async def get_meta(
    session: AsyncSession, document_id: uuid.UUID, *, refresh: bool = False
) -> FinancialStatementMeta | None:
    """Шапка формы документа (1:1) или `None`; `refresh` — перечитать атрибуты, даже если объект уже в сессии."""
    stmt = select(FinancialStatementMeta).where(FinancialStatementMeta.document_id == document_id)
    if refresh:
        stmt = stmt.execution_options(populate_existing=True)
    return (await session.execute(stmt)).scalar_one_or_none()


async def load_form(
    session: AsyncSession, meta_id: uuid.UUID
) -> tuple[list[FinancialStatementColumn], list[tuple[FinancialStatementLine, list[FinancialStatementValue]]]]:
    """Колонки по `column_index` и строки по `line_code` с их значениями по `column_index`."""
    columns = list(
        (
            await session.execute(
                select(FinancialStatementColumn)
                .where(FinancialStatementColumn.meta_id == meta_id)
                .order_by(FinancialStatementColumn.column_index)
            )
        ).scalars()
    )
    order = {column.id: column.column_index for column in columns}
    lines = list(
        (
            await session.execute(
                select(FinancialStatementLine)
                .where(FinancialStatementLine.meta_id == meta_id)
                .order_by(FinancialStatementLine.line_code, FinancialStatementLine.id)
            )
        ).scalars()
    )
    values = (
        await session.execute(
            select(FinancialStatementValue).where(FinancialStatementValue.line_id.in_([line.id for line in lines]))
        )
    ).scalars()
    by_line: dict[uuid.UUID, list[FinancialStatementValue]] = {line.id: [] for line in lines}
    for value in values:
        by_line[value.line_id].append(value)
    for items in by_line.values():
        items.sort(key=lambda value: order[value.column_id])
    return columns, [(line, by_line[line.id]) for line in lines]


async def _line_amount(
    session: AsyncSession,
    meta_id: uuid.UUID,
    line_code: str,
    *,
    current: bool = False,
    period_date: datetime.date | None = None,
) -> Decimal | None:
    """Сумма строки `line_code` в колонке `is_current` (`current`) или с датой `period_date`; нет — `None`."""
    stmt = (
        select(FinancialStatementValue.amount)
        .join(FinancialStatementLine, FinancialStatementLine.id == FinancialStatementValue.line_id)
        .join(FinancialStatementColumn, FinancialStatementColumn.id == FinancialStatementValue.column_id)
        .where(FinancialStatementLine.meta_id == meta_id, FinancialStatementLine.line_code == line_code)
        .order_by(FinancialStatementColumn.column_index)
        .limit(1)
    )
    if current:
        stmt = stmt.where(FinancialStatementColumn.is_current.is_(True))
    else:
        stmt = stmt.where(FinancialStatementColumn.period_date == period_date)
    return (await session.execute(stmt)).scalar_one_or_none()


async def _pair_document(session: AsyncSession, document: Document, pair_type: DocumentType) -> Document | None:
    """Форма-пара (D-5): другой тип формы того же ЮЛ, года и квартала, файл не `SUPERSEDED`; подтверждённая первой,
    затем самая поздняя по `created_at`. Документ без шапки формы не считается парой (иначе он затенял бы настоящую)."""
    stmt = (
        select(Document)
        .join(StoredFile, StoredFile.id == Document.stored_file_id)
        .join(FinancialStatementMeta, FinancialStatementMeta.document_id == Document.id)
        .where(
            Document.legal_entity_id == document.legal_entity_id,
            Document.period_year == document.period_year,
            Document.period_quarter == document.period_quarter,
            Document.document_type == pair_type,
            StoredFile.status != FileStatus.SUPERSEDED,
        )
        .order_by((Document.status == DocumentStatus.CONFIRMED).desc(), Document.created_at.desc())
        .limit(1)
    )
    return (await session.execute(stmt)).scalar_one_or_none()


async def pair_amounts(
    session: AsyncSession, document: Document
) -> tuple[tuple[Decimal | None, Decimal | None], Decimal | None]:
    """Вход межформенной проверки (D-5): `(1370 на отчётную дату, 1370 на 31.12 прошлого года)` из ББ и `2400` ОФР
    в колонке отчётного периода. Нет пары, шапки пары, строки или колонки — соответствующий `None`."""
    if document.period_year is None or document.period_quarter is None:
        return (None, None), None
    is_balance = document.document_type == DocumentType.BALANCE
    pair = await _pair_document(session, document, DocumentType.INCOME if is_balance else DocumentType.BALANCE)
    own_meta = await get_meta(session, document.id)
    pair_meta = await get_meta(session, pair.id) if pair is not None else None
    balance_meta, income_meta = (own_meta, pair_meta) if is_balance else (pair_meta, own_meta)
    balance: tuple[Decimal | None, Decimal | None] = (None, None)
    if balance_meta is not None:
        year_start = datetime.date(document.period_year - 1, 12, 31)
        balance = (
            await _line_amount(session, balance_meta.id, "1370", current=True),
            await _line_amount(session, balance_meta.id, "1370", period_date=year_start),
        )
    income = await _line_amount(session, income_meta.id, "2400", current=True) if income_meta is not None else None
    return balance, income