arrr-erp/app/modules/bank_statement_analyzer/service.py

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 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 / "Bank_Statement_Analysis.xlsx"
        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,
        }
        for path in paths:
            try:
                path.unlink(missing_ok=True)
            except OSError:
                pass
        input_dir = Path(job.job_directory) / "Input"
        try:
            input_dir.rmdir()
        except OSError:
            pass
        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 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(),
    }


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()