from __future__ import annotations import io import json import zipfile from dataclasses import dataclass from datetime import datetime, timezone from typing import Any from urllib.parse import quote from openpyxl import Workbook from openpyxl.worksheet.worksheet import Worksheet from sqlalchemy.ext.asyncio import AsyncSession from src.db.models.app_user import AppUser from src.db.models.form_type import FormTypeEnum from src.services.budget_form_service import BudgetFormService from src.services.export_normalizers import ( ExcelValueSerializer, Form1RowNormalizer, Form3ReportRowNormalizer, SmetaRowNormalizer, ) from src.services.export_writers import ( FORM1_FLAT_COLUMNS, FORM3_REPORT_COLUMNS, SMETA_COLUMNS, TabularSheetWriter, ) from src.services.project_service import ProjectService from src.services.sheet_service import SheetService XLSX_MEDIA_TYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" ZIP_MEDIA_TYPE = "application/zip" FORM1_DIRECTION_REQUIRED_SHEETS = {"AHR", "CAP"} SHEETS_WITH_SECTIONS = {"AHR", "CAP", "OPER"} @dataclass(frozen=True) class ExportSheetPlan: sheet_name: str effective_direction: str | None effective_sections: list[str] | None validation_direction: str | None validation_sections: list[str] | None class ExportService: def __init__(self, db: AsyncSession): self.budget_form_service = BudgetFormService(db) self.sheet_service = SheetService(db) self.project_service = ProjectService(db) self.value_serializer = ExcelValueSerializer() self.form1_row_normalizer = Form1RowNormalizer(self.value_serializer) self.form3_report_row_normalizer = Form3ReportRowNormalizer() self.smeta_row_normalizer = SmetaRowNormalizer() self.sheet_writers: dict[str, TabularSheetWriter] = { "__default__": TabularSheetWriter( columns=FORM1_FLAT_COLUMNS, row_normalizer=self.form1_row_normalizer.normalize, value_serializer=self.value_serializer.to_excel_value, ), "SMETA": TabularSheetWriter( columns=SMETA_COLUMNS, row_normalizer=self.smeta_row_normalizer.normalize, value_serializer=self.value_serializer.to_excel_value, ), } self.form3_report_writer = TabularSheetWriter( columns=FORM3_REPORT_COLUMNS, row_normalizer=self.form3_report_row_normalizer.normalize, value_serializer=self.value_serializer.to_excel_value, ) async def export_form_payload( self, form_id: int, current_user: AppUser, sheet: str | None = None, direction: str | None = None, sections: list[str] | None = None, ignore_non_applicable_query_params: bool = False, ) -> tuple[io.BytesIO, str, str]: stream, filename = await self.export_form_to_xlsx( form_id=form_id, current_user=current_user, sheet=sheet, direction=direction, sections=sections, ignore_non_applicable_query_params=ignore_non_applicable_query_params, ) return stream, filename, XLSX_MEDIA_TYPE async def export_project_report_payload( self, project_id: int, year: int, report_type: str, current_user: AppUser, sections: list[str] | None = None, ) -> tuple[io.BytesIO, str, str]: stream, filename = await self.export_project_report_to_xlsx( project_id=project_id, year=year, report_type=report_type, current_user=current_user, sections=sections, ) return stream, filename, XLSX_MEDIA_TYPE async def export_bulk_payload( self, form_ids: list[int], current_user: AppUser, skip_failed: bool = True, direction: str | None = None, sections: list[str] | None = None, ignore_non_applicable_query_params: bool = False, ) -> tuple[io.BytesIO, str, str]: archive_stream = io.BytesIO() exported_count = 0 errors: list[dict[str, Any]] = [] used_names: set[str] = set() with zipfile.ZipFile( archive_stream, mode="w", compression=zipfile.ZIP_DEFLATED, ) as archive: for form_id in self._deduplicate_ids(form_ids): try: form_stream, filename = await self.export_form_to_xlsx( form_id=form_id, current_user=current_user, direction=direction, sections=sections, ignore_non_applicable_query_params=ignore_non_applicable_query_params, ) archive_name = self._make_unique_filename(filename, used_names) archive.writestr(archive_name, form_stream.getvalue()) exported_count += 1 except ValueError as exc: if not skip_failed: raise errors.append({"form_id": form_id, "error": str(exc)}) if errors: archive.writestr( "errors.json", json.dumps({"errors": errors}, ensure_ascii=False, indent=2).encode("utf-8"), ) if exported_count == 0: raise ValueError("Не удалось экспортировать ни одной формы.") archive_stream.seek(0) filename = f"export_forms_{datetime.now(timezone.utc).strftime('%Y-%m-%d')}.zip" return archive_stream, filename, ZIP_MEDIA_TYPE async def export_form_to_xlsx( self, form_id: int, current_user: AppUser, sheet: str | None = None, direction: str | None = None, sections: list[str] | None = None, ignore_non_applicable_query_params: bool = False, ) -> tuple[io.BytesIO, str]: form = await self._get_form_for_export(form_id=form_id, current_user=current_user) sheets = self._resolve_requested_sheets(form=form, requested_sheet=sheet) self._validate_sections_for_form(form=form, sections=sections) plans = self._build_sheet_plans( form=form, sheets=sheets, direction=direction, sections=sections, ignore_non_applicable_query_params=ignore_non_applicable_query_params, ) workbook = self._create_workbook_with_metadata(form=form) exported_sheets = 0 for plan in plans: self._validate_sheet_query_params( form_type_code=form.form_type.code, sheet=plan.sheet_name, direction=plan.validation_direction, sections=plan.validation_sections, ) rows = await self.sheet_service.get( form_id=form_id, sheet=plan.sheet_name, direction=plan.effective_direction, sections=plan.effective_sections, ) worksheet = workbook.create_sheet(title=self._safe_sheet_name(plan.sheet_name)) self._write_sheet_rows( worksheet=worksheet, rows=rows, sheet_name=plan.sheet_name, ) exported_sheets += 1 if exported_sheets == 0: raise ValueError("Не удалось экспортировать листы формы.") stream = io.BytesIO() workbook.save(stream) workbook.close() stream.seek(0) filename = self._build_filename(form_id=form.id, form_type_code=form.form_type.code, year=form.year, sheet=sheet) return stream, filename async def export_project_report_to_xlsx( self, project_id: int, year: int, report_type: str, current_user: AppUser, sections: list[str] | None = None, ) -> tuple[io.BytesIO, str]: rows = await self.project_service.get_report_rows( project_id=project_id, year=year, report_type=report_type, sections=sections, user=current_user, ) report_type_upper = report_type.upper() workbook = Workbook() meta_sheet = workbook.active meta_sheet.title = "metadata" metadata_rows = [ ("Тип формы", "FORM_3"), ("Project ID", project_id), ("Year", year), ("Report Type", report_type_upper), ("Sections", ", ".join(sections) if sections else None), ("Экспортировано", datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")), ] for row_idx, (key, value) in enumerate(metadata_rows, start=1): meta_sheet.cell(row=row_idx, column=1, value=key) meta_sheet.cell(row=row_idx, column=2, value=value) worksheet = workbook.create_sheet(title=self._safe_sheet_name(report_type_upper)) self.form3_report_writer.write(worksheet, rows) stream = io.BytesIO() workbook.save(stream) workbook.close() stream.seek(0) date_str = datetime.now(timezone.utc).strftime("%Y-%m-%d") filename = f"{year}_FORM_3_{project_id}_{report_type_upper}_{date_str}.xlsx" return stream, filename @staticmethod def build_content_disposition(filename: str) -> str: encoded = quote(filename) return f"attachment; filename*=UTF-8''{encoded}" def _fill_metadata(self, sheet: Worksheet, form: Any) -> None: metadata_rows = [ ("Форма ID", form.id), ("Тип формы", form.form_type_code), ("Год", form.year), ("Org Unit ID", form.org_unit_id), ("Org Unit", form.org_unit.title if form.org_unit else None), ("Экспортировано", datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")), ] for row_idx, (key, value) in enumerate(metadata_rows, start=1): sheet.cell(row=row_idx, column=1, value=key) sheet.cell(row=row_idx, column=2, value=value) def _write_sheet_rows( self, worksheet: Worksheet, rows: list[tuple], sheet_name: str, ) -> None: writer = self.sheet_writers.get(sheet_name, self.sheet_writers["__default__"]) writer.write(worksheet, rows) async def _get_form_for_export(self, form_id: int, current_user: AppUser) -> Any: form = await self.budget_form_service.get( budget_form_id=form_id, user=current_user, load_form_type=True, load_org=True, ) if not form: raise ValueError(f"Форма {form_id} не найдена") return form def _resolve_requested_sheets(self, form: Any, requested_sheet: str | None) -> list[str]: if requested_sheet and requested_sheet not in form.form_type.sheet_list: raise ValueError(f"Лист {requested_sheet} формы {form.id} не найден") if requested_sheet: return [requested_sheet] return list(form.form_type.sheet_list) @staticmethod def _validate_sections_for_form(form: Any, sections: list[str] | None) -> None: if not sections: return invalid = sorted(set(sections) - set(form.form_type.section_list)) if invalid: raise ValueError(f"Недопустимые sections: {', '.join(invalid)}") def _create_workbook_with_metadata(self, form: Any) -> Workbook: workbook = Workbook() meta_sheet = workbook.active meta_sheet.title = "metadata" self._fill_metadata(meta_sheet, form) return workbook def _build_sheet_plans( self, form: Any, sheets: list[str], direction: str | None, sections: list[str] | None, ignore_non_applicable_query_params: bool, ) -> list[ExportSheetPlan]: plans: list[ExportSheetPlan] = [] for sheet_name in sheets: effective_direction = direction if self._requires_direction(form.form_type.code, sheet_name) else None effective_sections = sections if sheet_name in SHEETS_WITH_SECTIONS else None validation_direction = effective_direction if ignore_non_applicable_query_params else direction validation_sections = effective_sections if ignore_non_applicable_query_params else sections plans.append( ExportSheetPlan( sheet_name=sheet_name, effective_direction=effective_direction, effective_sections=effective_sections, validation_direction=validation_direction, validation_sections=validation_sections, ) ) return plans def _validate_sheet_query_params( self, form_type_code: FormTypeEnum, sheet: str, direction: str | None, sections: list[str] | None, ) -> None: if form_type_code == FormTypeEnum.FORM_1 and sheet in FORM1_DIRECTION_REQUIRED_SHEETS and not direction: raise ValueError("direction обязателен для FORM_1 листов AHR/CAP") if direction and not self._requires_direction(form_type_code, sheet): raise ValueError("direction допустим только для FORM_1 листов AHR/CAP") if sections and sheet not in SHEETS_WITH_SECTIONS: raise ValueError(f"sections не поддерживается для листа {sheet}") @staticmethod def _requires_direction(form_type_code: FormTypeEnum, sheet: str) -> bool: return form_type_code == FormTypeEnum.FORM_1 and sheet in FORM1_DIRECTION_REQUIRED_SHEETS @staticmethod def _safe_sheet_name(value: str) -> str: safe = value for char in ("[", "]", ":", "*", "?", "/", "\\"): safe = safe.replace(char, "_") safe = safe.strip() return (safe or "Sheet")[:31] @staticmethod def _build_filename(form_id: int, form_type_code: str, year: int, sheet: str | None = None) -> str: date_str = datetime.now(timezone.utc).strftime("%Y-%m-%d") if sheet: return f"{year}_{form_type_code}_{form_id}_{sheet}_{date_str}.xlsx" return f"{year}_{form_type_code}_{form_id}_{date_str}.xlsx" @staticmethod def _deduplicate_ids(form_ids: list[int]) -> list[int]: seen: set[int] = set() result: list[int] = [] for form_id in form_ids: if form_id not in seen: seen.add(form_id) result.append(form_id) return result @staticmethod def _make_unique_filename(filename: str, used_names: set[str]) -> str: if filename not in used_names: used_names.add(filename) return filename if "." in filename: base, ext = filename.rsplit(".", 1) ext = f".{ext}" else: base, ext = filename, "" suffix = 2 while True: candidate = f"{base}_{suffix}{ext}" if candidate not in used_names: used_names.add(candidate) return candidate suffix += 1