from __future__ import annotations import json import math import os import re import shutil import threading import time import uuid from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Iterable from fastapi import UploadFile from sqlalchemy import func, select, text from app.core.db.common import CommonSessionLocal from .analyzer import analyze_files, export_excel from .models import BankStatementAnalysisJob ALLOWED_ROLES = {"Partner", "Manager", "Branch Manager", "Staff", "Employee", "Consultant"} MAX_FILES = int(os.getenv("BANK_ANALYZER_MAX_FILES", "24")) MAX_FILE_BYTES = int(os.getenv("BANK_ANALYZER_MAX_FILE_MB", "50")) * 1024 * 1024 MAX_GLOBAL_PROCESSING = 3 MAX_PENDING_PER_USER = 3 COMPLETED_RETENTION_HOURS = 24 FAILED_RETENTION_HOURS = 24 POLL_SECONDS = 2 _SAFE = re.compile(r"[^A-Za-z0-9._-]+") _QUEUE_LOCK_ID = 820260713 _executor = ThreadPoolExecutor(max_workers=MAX_GLOBAL_PROCESSING, thread_name_prefix="bank-analyzer") _worker_lock = threading.Lock() _worker_started = False def _now() -> datetime: return datetime.now(timezone.utc) def _resolve_document_storage_root() -> Path: configured = (os.getenv("DOCUMENT_STORAGE_ROOT") or "").strip() if configured: configured_path = Path(configured).expanduser() return configured_path if configured_path.is_absolute() else (Path.cwd() / configured_path).resolve() if os.name == "nt": anchor = Path.cwd().anchor or (os.environ.get("SystemDrive", "C:") + "\\") return Path(anchor) / "AuditFirmERPDocuments" / "engagement_documents" return (Path.cwd() / "documents" / "engagement_documents").resolve() def _root() -> Path: return _resolve_document_storage_root().parent / "Work" def _segment(value: object, default: str = "NA") -> str: text_value = _SAFE.sub("_", str(value or default).strip()).strip("._-") return text_value[:80] or default def _workbook_filename(metas: list) -> str: """Build a safe BankName_ClientName.xlsx filename for every bank.""" bank_names = [str(getattr(meta, "bank_name", "") or "").strip() for meta in metas] client_names = [str(getattr(meta, "customer_name", "") or "").strip() for meta in metas] account_numbers = [str(getattr(meta, "account_number", "") or "").strip() for meta in metas] unique_banks = list(dict.fromkeys(name for name in bank_names if name)) unique_clients = list(dict.fromkeys(name for name in client_names if name)) unique_accounts = list(dict.fromkeys(name for name in account_numbers if name)) bank = unique_banks[0] if len(unique_banks) == 1 else ("Multiple_Banks" if unique_banks else "Bank") client = unique_clients[0] if len(unique_clients) == 1 else "" if not client: client = unique_accounts[0] if len(unique_accounts) == 1 else "Multiple_Statements" return f"{_segment(bank, 'Bank')}_{_segment(client, 'Client')}.xlsx" def role_bucket(roles: Iterable[str]) -> str | None: role_set = set(roles) for role, bucket in (("Partner", "Partner"), ("Manager", "Manager"), ("Branch Manager", "Manager"), ("Staff", "Staff"), ("Employee", "Staff"), ("Consultant", "Consultant")): if role in role_set: return bucket return None def can_use(roles: Iterable[str]) -> bool: return bool(set(roles) & ALLOWED_ROLES) def _user_root(user, roles: Iterable[str]) -> Path: bucket = role_bucket(roles) if not bucket: raise PermissionError("Bank Statement Analyzer is available only to Partner, Manager, Staff and Consultant roles.") label = _segment(getattr(user, "full_name", None) or getattr(user, "email", None) or f"user_{user.id}") return _root() / bucket / f"{int(user.id)}_{label}" / "Bank_Statement_Analyzer" def create_job_folder(user, roles: Iterable[str]) -> tuple[str, Path, Path]: job_id = uuid.uuid4().hex job_dir = _user_root(user, roles) / job_id input_dir = job_dir / "Input" output_dir = job_dir / "Output" input_dir.mkdir(parents=True, exist_ok=False) output_dir.mkdir(parents=True, exist_ok=False) return job_id, input_dir, output_dir def _validate_pdf_header(data: bytes) -> None: if not data.startswith(b"%PDF-"): raise ValueError("Only genuine PDF files are allowed.") async def save_uploads(files: list[UploadFile], input_dir: Path) -> list[Path]: usable = [item for item in files if item and (item.filename or "").strip()] if not usable: raise ValueError("Please select at least one PDF bank statement.") if len(usable) > MAX_FILES: raise ValueError(f"A maximum of {MAX_FILES} PDF files can be analyzed in one job.") saved: list[Path] = [] try: for index, upload in enumerate(usable, start=1): name = Path(upload.filename or f"statement_{index}.pdf").name if Path(name).suffix.lower() != ".pdf": raise ValueError(f"{name}: only PDF files are allowed.") target = input_dir / f"{index:02d}_{_segment(Path(name).stem, f'statement_{index}')}.pdf" size = 0 first = b"" with target.open("wb") as handle: while True: chunk = await upload.read(1024 * 1024) if not chunk: break if not first: first = chunk[:8] size += len(chunk) if size > MAX_FILE_BYTES: raise ValueError(f"{name}: file exceeds the {MAX_FILE_BYTES // (1024 * 1024)} MB limit.") handle.write(chunk) _validate_pdf_header(first) saved.append(target) return saved except Exception: shutil.rmtree(input_dir.parent, ignore_errors=True) raise def pending_count_for_user(user_id: int) -> int: db = CommonSessionLocal() try: return int(db.scalar(select(func.count()).select_from(BankStatementAnalysisJob).where(BankStatementAnalysisJob.user_id == int(user_id), BankStatementAnalysisJob.status.in_(["queued", "processing"]))) or 0) finally: db.close() def enqueue_job(*, user, roles: Iterable[str], job_id: str, paths: list[Path], job_dir: Path, bank_selection: str, financial_year: str, customer_override: str, account_override: str, classification_enabled: bool) -> BankStatementAnalysisJob: if pending_count_for_user(int(user.id)) >= MAX_PENDING_PER_USER: shutil.rmtree(job_dir, ignore_errors=True) raise ValueError("You already have three queued or processing analyses. Please wait for one to complete before submitting another.") bucket = role_bucket(roles) if not bucket: shutil.rmtree(job_dir, ignore_errors=True) raise PermissionError("You do not have access to the Bank Statement Analyzer.") job = BankStatementAnalysisJob( id=job_id, tenant_id=getattr(user, "tenant_id", None), branch_id=getattr(user, "branch_id", None), user_id=int(user.id), role_bucket=bucket, selected_bank=bank_selection, financial_year=(financial_year or "").strip() or None, customer_override=(customer_override or "").strip() or None, account_override=(account_override or "").strip() or None, classification_enabled=bool(classification_enabled), status="queued", progress_percent=0, file_count=len(paths), input_files_json=json.dumps([str(path) for path in paths]), job_directory=str(job_dir), submitted_at_utc=_now(), ) db = CommonSessionLocal() try: db.add(job) db.commit() db.refresh(job) except Exception: db.rollback() shutil.rmtree(job_dir, ignore_errors=True) raise finally: db.close() ensure_worker_started() return job def _claim_jobs() -> list[str]: db = CommonSessionLocal() try: try: db.execute(text("SELECT pg_advisory_xact_lock(:lock_id)"), {"lock_id": _QUEUE_LOCK_ID}) except Exception: db.rollback() processing = int(db.scalar(select(func.count()).select_from(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "processing")) or 0) capacity = max(0, MAX_GLOBAL_PROCESSING - processing) if capacity <= 0: return [] stmt = select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "queued").order_by(BankStatementAnalysisJob.submitted_at_utc.asc()).limit(capacity) try: stmt = stmt.with_for_update(skip_locked=True) except Exception: pass jobs = list(db.scalars(stmt).all()) now = _now() for job in jobs: job.status = "processing" job.progress_percent = 5 job.started_at_utc = now job.error_message = None db.commit() return [job.id for job in jobs] except Exception: db.rollback() return [] finally: db.close() def _process_job(job_id: str) -> None: db = CommonSessionLocal() try: job = db.get(BankStatementAnalysisJob, job_id) if not job or job.status != "processing": return paths = [Path(value) for value in json.loads(job.input_files_json)] output_dir = Path(job.job_directory) / "Output" job.progress_percent = 15 db.commit() metas, all_df, unique_df = analyze_files( paths, job.customer_override or "", job.account_override or "", bank_hint=job.selected_bank, classification_enabled=job.classification_enabled, ) job.progress_percent = 75 db.commit() output = output_dir / _workbook_filename(metas) export_excel( output, metas, all_df, unique_df, financial_year=job.financial_year or "", selected_bank=job.selected_bank, classification_enabled=job.classification_enabled, ) summary = { "job_id": job.id, "statement_count": len(metas), "rows_extracted": int(len(all_df)), "unique_transactions": int(len(unique_df)), "exact_duplicate_rows": int(all_df["exact_duplicate"].fillna(False).astype(bool).sum()) if not all_df.empty and "exact_duplicate" in all_df.columns else 0, "possible_duplicate_rows": int(all_df["possible_duplicate"].fillna(False).astype(bool).sum()) if not all_df.empty and "possible_duplicate" in all_df.columns else 0, "review_items": int(unique_df.review_note.fillna("").ne("").sum()) if not unique_df.empty and "review_note" in unique_df.columns else 0, "categories": int(unique_df.category.nunique()) if not unique_df.empty and "category" in unique_df.columns else 0, "banks": sorted({meta.bank_name for meta in metas}), "customer_name": next((meta.customer_name for meta in metas if meta.customer_name), ""), "account_number": next((meta.account_number for meta in metas if meta.account_number), ""), "financial_year": job.financial_year or "", "classification_enabled": job.classification_enabled, } # Keep the original uploaded statements until the job expiry time. # This applies equally to completed and failed jobs and allows the # uploader to reproduce parser issues during the 24-hour retention window. job.output_file = str(output) job.summary_json = json.dumps(summary) job.status = "completed" job.progress_percent = 100 job.completed_at_utc = _now() job.expires_at_utc = _now() + timedelta(hours=COMPLETED_RETENTION_HOURS) db.commit() except Exception as exc: db.rollback() job = db.get(BankStatementAnalysisJob, job_id) if job: job.status = "failed" job.progress_percent = 100 job.error_message = str(exc)[:4000] job.completed_at_utc = _now() job.expires_at_utc = _now() + timedelta(hours=FAILED_RETENTION_HOURS) db.commit() finally: db.close() def _cleanup_expired() -> None: db = CommonSessionLocal() try: jobs = list(db.scalars(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.expires_at_utc.is_not(None), BankStatementAnalysisJob.expires_at_utc < _now())).all()) for job in jobs: shutil.rmtree(job.job_directory, ignore_errors=True) db.delete(job) db.commit() except Exception: db.rollback() finally: db.close() def _worker_loop() -> None: while True: try: _cleanup_expired() for job_id in _claim_jobs(): _executor.submit(_process_job, job_id) except Exception: pass time.sleep(POLL_SECONDS) def ensure_worker_started() -> None: global _worker_started with _worker_lock: if _worker_started: return thread = threading.Thread(target=_worker_loop, name="bank-analyzer-queue", daemon=True) thread.start() _worker_started = True def _summary(job: BankStatementAnalysisJob) -> dict: return json.loads(job.summary_json) if job.summary_json else {} def average_duration_minutes() -> float: db = CommonSessionLocal() try: jobs = list(db.scalars(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "completed", BankStatementAnalysisJob.started_at_utc.is_not(None), BankStatementAnalysisJob.completed_at_utc.is_not(None)).order_by(BankStatementAnalysisJob.completed_at_utc.desc()).limit(20)).all()) durations = [(job.completed_at_utc - job.started_at_utc).total_seconds() / 60 for job in jobs if job.completed_at_utc and job.started_at_utc] return max(1.0, sum(durations) / len(durations)) if durations else 4.0 finally: db.close() def queue_position(job: BankStatementAnalysisJob) -> int | None: if job.status != "queued": return None db = CommonSessionLocal() try: return int(db.scalar(select(func.count()).select_from(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "queued", BankStatementAnalysisJob.submitted_at_utc <= job.submitted_at_utc)) or 1) finally: db.close() def estimated_wait(job: BankStatementAnalysisJob) -> str: if job.status == "processing": return "Processing now" if job.status != "queued": return "" position = queue_position(job) or 1 waves = max(1, math.ceil(position / MAX_GLOBAL_PROCESSING)) average = average_duration_minutes() low = max(1, math.ceil((waves - 1) * average)) high = max(5, math.ceil(waves * average + 2)) return f"Approximately {low}-{high} minutes" def get_owned_job(user_id: int, job_id: str) -> BankStatementAnalysisJob | None: if not re.fullmatch(r"[a-f0-9]{32}", job_id or ""): return None db = CommonSessionLocal() try: job = db.scalar(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.id == job_id, BankStatementAnalysisJob.user_id == int(user_id))) if job: db.expunge(job) return job finally: db.close() def list_user_jobs(user_id: int, limit: int = 25) -> list[BankStatementAnalysisJob]: db = CommonSessionLocal() try: jobs = list(db.scalars(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.user_id == int(user_id)).order_by(BankStatementAnalysisJob.submitted_at_utc.desc()).limit(limit)).all()) for job in jobs: db.expunge(job) return jobs finally: db.close() def original_statement_files(job: BankStatementAnalysisJob) -> list[Path]: """Return retained input PDFs that still belong to the job directory.""" job_root = Path(job.job_directory).resolve() input_root = (job_root / "Input").resolve() try: configured = [Path(value).resolve() for value in json.loads(job.input_files_json or "[]")] except (TypeError, ValueError, json.JSONDecodeError): configured = [] retained: list[Path] = [] for path in configured: try: path.relative_to(input_root) except ValueError: continue if path.is_file() and path.suffix.lower() == ".pdf": retained.append(path) return retained def get_owned_original_statement(user_id: int, job_id: str, file_index: int) -> tuple[BankStatementAnalysisJob, Path] | None: job = get_owned_job(user_id, job_id) if not job or job.status not in {"completed", "failed"}: return None files = original_statement_files(job) if file_index < 1 or file_index > len(files): return None return job, files[file_index - 1] def job_view(job: BankStatementAnalysisJob) -> dict: return { "id": job.id, "status": job.status, "progress_percent": job.progress_percent, "file_count": job.file_count, "selected_bank": job.selected_bank, "financial_year": job.financial_year or "", "submitted_at": job.submitted_at_utc, "completed_at": job.completed_at_utc, "expires_at": job.expires_at_utc, "error_message": job.error_message or "", "queue_position": queue_position(job), "estimated_wait": estimated_wait(job), "summary": _summary(job), "download_ready": job.status == "completed" and bool(job.output_file) and Path(job.output_file).is_file(), "original_files": [ {"index": index, "name": path.name} for index, path in enumerate(original_statement_files(job), start=1) ], "originals_retained": job.status in {"completed", "failed"} and bool(original_statement_files(job)), } def delete_owned_job(user_id: int, job_id: str) -> bool: db = CommonSessionLocal() try: job = db.scalar(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.id == job_id, BankStatementAnalysisJob.user_id == int(user_id))) if not job: return False if job.status == "processing": raise ValueError("A processing analysis cannot be deleted. Please wait for it to finish.") shutil.rmtree(job.job_directory, ignore_errors=True) db.delete(job) db.commit() return True finally: db.close()