nifodea/infrastructure/repositories/excel_repository.py

"""ExcelRepository — хранит FinancialModel в .xlsx (5 листов + ModelInfo)."""
from pathlib import Path
from uuid import UUID

from openpyxl import Workbook, load_workbook

from application.repositories.model_repository import ModelRepository
from domain import (
    Account,
    Asset,
    ExchangeRate,
    FinancialModel,
    Liability,
    RecurringCashflow,
    Transaction,
)


class ExcelRepositoryError(Exception):
    pass


_SHEET_CONFIG = {
    "Accounts": {
        "fields": ["id", "name", "currency", "balance"],
        "cls": Account,
    },
    "Transactions": {
        "fields": ["id", "date", "account", "category", "amount", "description"],
        "cls": Transaction,
    },
    "Recurring": {
        "fields": ["id", "start_date", "end_date", "frequency", "amount", "category"],
        "cls": RecurringCashflow,
    },
    "Assets": {
        "fields": ["id", "name", "value", "growth_rate"],
        "cls": Asset,
    },
    "Liabilities": {
        "fields": ["id", "name", "balance", "interest", "payment"],
        "cls": Liability,
    },
    "ExchangeRates": {
        "fields": ["from_currency", "to_currency", "rate"],
        "cls": ExchangeRate,
    },
}


class ExcelRepository(ModelRepository):
    def load(self, path: str | Path) -> FinancialModel:
        p = Path(path)
        if not p.exists():
            raise ExcelRepositoryError(f"File not found: {p}")

        wb = load_workbook(p, read_only=True, data_only=True)
        model = FinancialModel()
        model.exchange_rates = []

        for sheet_name, config in _SHEET_CONFIG.items():
            if sheet_name not in wb.sheetnames:
                continue
            ws = wb[sheet_name]
            rows = list(ws.iter_rows(values_only=True))
            if len(rows) < 2:
                continue

            headers = [str(h).strip().lower() if h else "" for h in rows[0]]
            for row in rows[1:]:
                if not any(v is not None for v in row):
                    continue
                data = {}
                for i, header in enumerate(headers):
                    val = row[i] if i < len(row) else None
                    if val is not None:
                        data[header] = str(val) if not isinstance(val, (int, float)) else val
                self._add_to_model(model, sheet_name, data, config["cls"])

        if "ModelInfo" in wb.sheetnames:
            ws = wb["ModelInfo"]
            rows = list(ws.iter_rows(values_only=True))
            for row in rows:
                if row[0] and str(row[0]).strip().lower() == "base_currency" and len(row) > 1:
                    model.base_currency = str(row[1]).strip()

        wb.close()
        return model

    def save(self, model: FinancialModel, path: str | Path) -> None:
        p = Path(path)
        wb = Workbook()
        wb.remove(wb.active)

        collections = {
            "Accounts": model.accounts,
            "Transactions": model.transactions,
            "Recurring": model.recurring,
            "Assets": model.assets,
            "Liabilities": model.liabilities,
            "ExchangeRates": model.exchange_rates,
        }

        for sheet_name, items in collections.items():
            config = _SHEET_CONFIG[sheet_name]
            ws = wb.create_sheet(title=sheet_name)
            ws.append(config["fields"])
            for item in items:
                row = [
                    str(getattr(item, f)) if isinstance(getattr(item, f), UUID)
                    else getattr(item, f)
                    for f in config["fields"]
                ]
                ws.append(row)

        ws_info = wb.create_sheet(title="ModelInfo")
        ws_info.append(["Property", "Value"])
        ws_info.append(["base_currency", model.base_currency])

        wb.save(p)

    def _add_to_model(self, model: FinancialModel, sheet_name: str, data: dict, cls) -> None:
        try:
            instance = cls.model_validate(data)
        except Exception as e:
            raise ExcelRepositoryError(f"Failed to parse row in {sheet_name}: {e}") from e
        if sheet_name == "Accounts":
            model.accounts.append(instance)
        elif sheet_name == "Transactions":
            model.transactions.append(instance)
        elif sheet_name == "Recurring":
            model.recurring.append(instance)
        elif sheet_name == "Assets":
            model.assets.append(instance)
        elif sheet_name == "Liabilities":
            model.liabilities.append(instance)
        elif sheet_name == "ExchangeRates":
            model.exchange_rates.append(instance)