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

from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

from fastapi import APIRouter, File, Form, Request, UploadFile
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
from sqlalchemy import select

from app.core.db.common import CommonSessionLocal
from app.core.http_responses import ui_access_denied, not_found_response
from app.core.security.csrf import get_or_create_csrf_token, validate_csrf
from app.core.security.session_auth import get_current_user
from app.core.templating import templates
from app.modules.core.rbac.deps import get_user_permissions, get_user_roles

from .parsers.registry import BANK_OPTIONS
from .service import (
    can_use,
    create_job_folder,
    delete_owned_job,
    enqueue_job,
    ensure_worker_started,
    get_owned_job,
    get_owned_original_statement,
    job_view,
    list_user_jobs,
    save_uploads,
)

router = APIRouter(prefix="/tools/bank-statement-analyzer", tags=["bank-statement-analyzer-ui"])


def _display_timezone_name(request: Request, db, user) -> str:
    """Resolve the active branch timezone, falling back to tenant and India defaults."""
    from app.modules.core.tenancy.models import Branch, Tenant

    branch_id = request.session.get("active_branch_id") or getattr(user, "branch_id", None)
    tenant_id = request.session.get("active_tenant_id") or getattr(user, "tenant_id", None)

    if branch_id not in (None, "", 0, "0"):
        branch = db.scalar(select(Branch).where(Branch.id == int(branch_id)))
        branch_timezone = (getattr(branch, "timezone", None) or "").strip() if branch else ""
        if branch_timezone:
            return branch_timezone

    if tenant_id not in (None, "", 0, "0"):
        tenant = db.scalar(select(Tenant).where(Tenant.id == int(tenant_id)))
        tenant_timezone = (getattr(tenant, "default_timezone", None) or "").strip() if tenant else ""
        if tenant_timezone:
            return tenant_timezone

    return "Asia/Kolkata"


def _format_local_datetime(value: datetime | None, timezone_name: str) -> str | None:
    if value is None:
        return None
    try:
        target_timezone = ZoneInfo(timezone_name)
    except ZoneInfoNotFoundError:
        target_timezone = ZoneInfo("Asia/Kolkata")
    aware_value = value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
    return aware_value.astimezone(target_timezone).strftime("%d-%b-%Y %I:%M %p %Z")


def _localised_job_view(job, timezone_name: str) -> dict:
    view = job_view(job)
    view["submitted_at"] = _format_local_datetime(view.get("submitted_at"), timezone_name)
    view["completed_at"] = _format_local_datetime(view.get("completed_at"), timezone_name)
    view["expires_at"] = _format_local_datetime(view.get("expires_at"), timezone_name)
    return view


def _ctx(request, db, user, **extra):
    data = {
        "request": request,
        "current_user": user,
        "current_user_roles": get_user_roles(db, user.id),
        "current_user_permissions": get_user_permissions(db, user.id),
        "csrf_token": get_or_create_csrf_token(request),
        "title": "Bank Statement Analyzer",
        "bank_options": BANK_OPTIONS,
        "selected_bank": "auto",
        "financial_year": "",
        "classification_enabled": True,
    }
    data.update(extra)
    return data


def _auth(request, db):
    user = get_current_user(request, db=db)
    if not user:
        return None, None, RedirectResponse("/login", status_code=303)
    roles = get_user_roles(db, user.id)
    if not can_use(roles):
        return user, roles, ui_access_denied("Bank Statement Analyzer is available only to Partner, Manager, Staff and Consultant roles.")
    return user, roles, None


@router.get("")
def index(request: Request, job: str | None = None):
    ensure_worker_started()
    db = CommonSessionLocal()
    try:
        user, roles, denied = _auth(request, db)
        if denied:
            return denied
        timezone_name = _display_timezone_name(request, db, user)
        selected_job = get_owned_job(user.id, job) if job else None
        recent = [_localised_job_view(item, timezone_name) for item in list_user_jobs(user.id, limit=8)]
        active_job = _localised_job_view(selected_job, timezone_name) if selected_job else None
        return templates.TemplateResponse("modules/bank_statement_analyzer/templates/bank_statement_analyzer/index.html", _ctx(request, db, user, error="", active_job=active_job, recent_jobs=recent, display_timezone=timezone_name))
    finally:
        db.close()


@router.post("/analyze")
async def analyze(request: Request, csrf_token: str = Form(...), bank_selection: str = Form("auto"), financial_year: str = Form(""), customer_name: str = Form(""), account_number: str = Form(""), enable_classification: str | None = Form(None), statements: list[UploadFile] = File(...)):
    db = CommonSessionLocal()
    selected_bank = bank_selection if bank_selection in dict(BANK_OPTIONS) else "auto"
    classification_enabled = enable_classification == "1"
    job_dir: Path | None = None
    try:
        user, roles, denied = _auth(request, db)
        if denied:
            return denied
        validate_csrf(request, csrf_token)
        job_id, input_dir, _output_dir = create_job_folder(user, roles)
        job_dir = input_dir.parent
        paths = await save_uploads(statements, input_dir)
        enqueue_job(user=user, roles=roles, job_id=job_id, paths=paths, job_dir=job_dir, bank_selection=selected_bank, financial_year=financial_year, customer_override=customer_name, account_override=account_number, classification_enabled=classification_enabled)
        return RedirectResponse(f"/tools/bank-statement-analyzer?job={job_id}#analysis-status", status_code=303)
    except Exception as exc:
        if job_dir:
            import shutil
            shutil.rmtree(job_dir, ignore_errors=True)
        user = get_current_user(request, db=db)
        if not user:
            return RedirectResponse("/login", status_code=303)
        timezone_name = _display_timezone_name(request, db, user)
        recent = [_localised_job_view(item, timezone_name) for item in list_user_jobs(user.id, limit=8)]
        return templates.TemplateResponse("modules/bank_statement_analyzer/templates/bank_statement_analyzer/index.html", _ctx(request, db, user, error=str(exc), active_job=None, recent_jobs=recent, selected_bank=selected_bank, financial_year=financial_year, classification_enabled=classification_enabled, display_timezone=timezone_name), status_code=400)
    finally:
        db.close()


@router.get("/jobs")
def jobs(request: Request):
    ensure_worker_started()
    db = CommonSessionLocal()
    try:
        user, roles, denied = _auth(request, db)
        if denied:
            return denied
        timezone_name = _display_timezone_name(request, db, user)
        items = [_localised_job_view(item, timezone_name) for item in list_user_jobs(user.id, limit=100)]
        return templates.TemplateResponse("modules/bank_statement_analyzer/templates/bank_statement_analyzer/jobs.html", _ctx(request, db, user, jobs=items, display_timezone=timezone_name))
    finally:
        db.close()


@router.get("/jobs/{job_id}/status")
def status(job_id: str, request: Request):
    ensure_worker_started()
    db = CommonSessionLocal()
    try:
        user, roles, denied = _auth(request, db)
        if denied:
            return JSONResponse({"detail": "Access denied"}, status_code=403)
        job = get_owned_job(user.id, job_id)
        if not job:
            return JSONResponse({"detail": "Analysis job not found"}, status_code=404)
        view = job_view(job)
        return JSONResponse({"id": view["id"], "status": view["status"], "progress_percent": view["progress_percent"], "queue_position": view["queue_position"], "estimated_wait": view["estimated_wait"], "download_ready": view["download_ready"], "error_message": view["error_message"]})
    finally:
        db.close()


@router.get("/jobs/{job_id}/download")
def download(job_id: str, request: Request):
    db = CommonSessionLocal()
    try:
        user, roles, denied = _auth(request, db)
        if denied:
            return denied
        job = get_owned_job(user.id, job_id)
        if not job:
            return not_found_response(request, "Analysis job not found or expired.")
        view = job_view(job)
        if not view["download_ready"]:
            return not_found_response(request, "Analysis workbook is not ready or has expired.")
        return FileResponse(path=job.output_file, filename=Path(job.output_file).name, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
    finally:
        db.close()


@router.get("/jobs/{job_id}/statements/{file_index}/download")
def download_original_statement(job_id: str, file_index: int, request: Request):
    db = CommonSessionLocal()
    try:
        user, roles, denied = _auth(request, db)
        if denied:
            return denied
        resolved = get_owned_original_statement(user.id, job_id, file_index)
        if not resolved:
            return not_found_response(request, "Original statement is unavailable or has expired.")
        _job, path = resolved
        return FileResponse(
            path=path,
            filename=path.name,
            media_type="application/pdf",
        )
    finally:
        db.close()


@router.post("/jobs/{job_id}/delete")
def delete(job_id: str, request: Request, csrf_token: str = Form(...)):
    db = CommonSessionLocal()
    try:
        user, roles, denied = _auth(request, db)
        if denied:
            return denied
        validate_csrf(request, csrf_token)
        delete_owned_job(user.id, job_id)
        return RedirectResponse("/tools/bank-statement-analyzer/jobs", status_code=303)
    finally:
        db.close()