finance/backend/app/categorization/service.py

import uuid
from collections.abc import Iterable

from sqlalchemy.orm import Session

from app.categorization import mcc
from app.categorization.models import (
    CategorizationRule,
    CategoryMcc,
    IncomeSource,
    MccCode,
    Payee,
    PayeeAlias,
)
from app.errors import conflict, not_found
from app.ledger.models import Category


def resolve_category_name(category: Category) -> str | None:
    """Literal display name, or None when the name localizes via the system key.

    System categories serialize as a pure ``nameKey`` (for example ``savings``),
    never as a ``category.*`` prefixed string.
    """
    if category.custom_name is not None:
        return category.custom_name
    if category.name_override is not None:
        return category.name_override
    return None


def resolve_category_name_key(category: Category) -> str | None:
    """Pure localization key for system categories; None for user categories."""
    return category.system_key


def get_income_source(db: Session, *, space_id, income_source_id) -> IncomeSource:
    obj = db.query(IncomeSource).filter_by(id=income_source_id, space_id=space_id).one_or_none()
    if obj is None:
        raise not_found("INCOME_SOURCE_NOT_FOUND", "Income source not found")
    return obj


def ensure_income_source(db: Session, *, space_id, name: str) -> IncomeSource:
    obj = db.query(IncomeSource).filter_by(space_id=space_id, name=name).one_or_none()
    if obj is None:
        obj = IncomeSource(space_id=space_id, name=name)
        db.add(obj)
        db.flush()
    return obj


def get_or_create_payee(db: Session, *, space_id, name: str, alias: str | None = None) -> Payee:
    payee = db.query(Payee).filter_by(space_id=space_id, name=name).one_or_none()
    if payee is None:
        payee = Payee(space_id=space_id, name=name)
        db.add(payee)
        db.flush()
    if alias:
        add_alias(db, space_id=space_id, payee_id=payee.id, alias=alias)
    return payee


def add_alias(db: Session, *, space_id, payee_id, alias: str) -> PayeeAlias:
    existing = db.query(PayeeAlias).filter_by(space_id=space_id, alias=alias).one_or_none()
    if existing is not None:
        if existing.payee_id != payee_id:
            raise conflict(
                "PAYEE_ALIAS_TAKEN",
                "Alias already mapped to another payee",
            )
        return existing
    obj = PayeeAlias(space_id=space_id, payee_id=payee_id, alias=alias)
    db.add(obj)
    db.flush()
    return obj


def find_payee_by_description(db: Session, *, space_id, description: str) -> Payee | None:
    if not description:
        return None
    alias = (
        db.query(PayeeAlias)
        .filter(PayeeAlias.space_id == space_id, PayeeAlias.alias == description)
        .one_or_none()
    )
    if alias is not None:
        return db.get(Payee, alias.payee_id)
    payee = db.query(Payee).filter_by(space_id=space_id, name=description).one_or_none()
    return payee


def rule_category_for_payee(db: Session, *, space_id, payee_id) -> CategorizationRule | None:
    return (
        db.query(CategorizationRule).filter_by(space_id=space_id, payee_id=payee_id).one_or_none()
    )


def set_payee_rule(db: Session, *, space_id, payee_id, category_id, user_id) -> CategorizationRule:
    payee = db.query(Payee).filter_by(id=payee_id, space_id=space_id).one_or_none()
    if payee is None:
        raise not_found("PAYEE_NOT_FOUND", "Payee not found")
    rule = rule_category_for_payee(db, space_id=space_id, payee_id=payee_id)
    if rule is None:
        rule = CategorizationRule(
            space_id=space_id,
            payee_id=payee_id,
            category_id=category_id,
            created_by_user_id=user_id,
        )
        db.add(rule)
    else:
        rule.category_id = category_id
    db.flush()
    return rule


def ensure_mcc_dictionary(db: Session) -> None:
    """Idempotently seed the global ``mcc_codes`` table from iso18245."""
    existing = {code for (code,) in db.query(MccCode.code).all()}
    for code, title in mcc.iter_mcc_dictionary():
        if code in existing:
            continue
        db.add(MccCode(code=code, title=title[:300], dictionary_version=mcc.MCC_DICTIONARY_VERSION))
    db.flush()


def ensure_others_category(db: Session, *, space_id: uuid.UUID) -> Category:
    """The system «Остальное» fallback category, created on first use."""
    others = (
        db.query(Category)
        .filter_by(space_id=space_id, system_key=mcc.OTHERS_SYSTEM_KEY)
        .one_or_none()
    )
    if others is None:
        others = Category(
            space_id=space_id,
            kind="monthly",
            system_key=mcc.OTHERS_SYSTEM_KEY,
            custom_name=mcc.OTHERS_CATEGORY_NAME,
        )
        db.add(others)
        db.flush()
    return others


def ensure_default_categories(db: Session, *, space_id: uuid.UUID) -> dict[str, Category]:
    """Idempotently pre-create the 20 popular categories plus «Остальное».

    Each default category is keyed by ``system_key`` and carries its MCC code
    links; every link is only created while its MCC code is unclaimed inside
    the space, so user re-mappings are never overwritten.
    """
    ensure_mcc_dictionary(db)
    keys = mcc.default_category_system_keys()
    existing: dict[str, Category] = {
        category.system_key: category
        for category in db.query(Category)
        .filter(Category.space_id == space_id, Category.system_key.in_(keys))
        .all()
        if category.system_key is not None
    }
    result: dict[str, Category] = {}
    claimed = {
        code for (code,) in db.query(CategoryMcc.mcc_code).filter_by(space_id=space_id).all()
    }
    for definition in mcc.default_categories():
        category = existing.get(definition.system_key)
        if category is None:
            category = Category(
                space_id=space_id,
                kind="monthly",
                system_key=definition.system_key,
                custom_name=definition.name,
            )
            db.add(category)
            db.flush()
        for code in definition.mcc_codes:
            if code in claimed:
                continue
            db.add(CategoryMcc(space_id=space_id, category_id=category.id, mcc_code=code))
            claimed.add(code)
        result[definition.system_key] = category
    result[mcc.OTHERS_SYSTEM_KEY] = ensure_others_category(db, space_id=space_id)
    db.flush()
    return result


def category_for_mcc(db: Session, *, space_id: uuid.UUID, code: int) -> Category | None:
    """Category owning the MCC code, else the system «Остальное» fallback."""
    link = db.query(CategoryMcc).filter_by(space_id=space_id, mcc_code=code).one_or_none()
    if link is not None:
        category = db.get(Category, link.category_id)
        if category is not None and not category.archived:
            return category
    return (
        db.query(Category)
        .filter_by(space_id=space_id, system_key=mcc.OTHERS_SYSTEM_KEY)
        .one_or_none()
    )


def get_category_mcc_codes(
    db: Session, *, space_id: uuid.UUID, category_id: uuid.UUID
) -> list[int]:
    codes = (
        db.query(CategoryMcc.mcc_code)
        .filter_by(space_id=space_id, category_id=category_id)
        .order_by(CategoryMcc.mcc_code)
        .all()
    )
    return [code for (code,) in codes]


def set_category_mcc_codes(
    db: Session,
    *,
    space_id: uuid.UUID,
    category_id: uuid.UUID,
    codes: Iterable[int],
) -> list[int]:
    """Replace the category's MCC code set; codes are unique per space."""
    category = db.query(Category).filter_by(id=category_id, space_id=space_id).one_or_none()
    if category is None:
        raise not_found("BUDGET_CATEGORY_NOT_FOUND", "Category not found")
    wanted = sorted({int(code) for code in codes})
    if wanted and category.archived:
        raise conflict("BUDGET_CATEGORY_ARCHIVED", "Archived categories cannot have MCC codes")
    if wanted and category.kind != "monthly":
        raise conflict("CATEGORY_MCC_KIND_INVALID", "Only monthly categories can have MCC codes")
    for code in wanted:
        if db.get(MccCode, code) is None:
            raise not_found("MCC_CODE_NOT_FOUND", f"Unknown MCC code {code}")
    current = {
        row for row in db.query(CategoryMcc).filter_by(space_id=space_id, category_id=category_id)
    }
    wanted_set = set(wanted)
    for row in current:
        if row.mcc_code not in wanted_set:
            db.delete(row)
    owned_elsewhere = {
        code
        for (code,) in (
            db.query(CategoryMcc.mcc_code)
            .filter(
                CategoryMcc.space_id == space_id,
                CategoryMcc.category_id != category_id,
                CategoryMcc.mcc_code.in_(wanted_set),
            )
            .all()
        )
    }
    if owned_elsewhere:
        raise conflict(
            "CATEGORY_MCC_TAKEN",
            "MCC codes already mapped to another category",
            details={"mcc_codes": sorted(owned_elsewhere)},
        )
    current_codes = {row.mcc_code for row in current}
    for code in wanted:
        if code not in current_codes:
            db.add(CategoryMcc(space_id=space_id, category_id=category_id, mcc_code=code))
    db.flush()
    return wanted