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_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 others = existing.get(mcc.OTHERS_SYSTEM_KEY) 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() result[mcc.OTHERS_SYSTEM_KEY] = others 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