from __future__ import annotations

from datetime import date, datetime, timezone
import asyncio
import logging
from pathlib import Path
from urllib.parse import quote
import re

from fastapi import APIRouter, File, Form, Header, Request, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse, Response
from sqlalchemy import select
from sqlalchemy.orm import joinedload

from app.core.db.common import CommonSessionLocal
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 app.modules.core.rbac.permission_guard import require_permission
from app.modules.documents.services import (
    DOCUMENT_TYPES,
    build_document_scope,
    get_document,
    get_download_request,
    get_download_request_for_node,
    get_latest_version,
    get_version,
    acknowledge_storage_job,
    cleanup_completed_storage_job_vps_stage,
    cleanup_completed_permanent_storage_job_vps_stage,
    retry_verified_vps_cleanup_for_node,
    authenticate_storage_node,
    create_branch_storage_node,
    generate_storage_secret,
    hash_storage_secret,
    get_storage_job_for_node,
    list_documents_for_engagement,
    list_download_requests,
    list_udin_records,
    list_recent_download_requests_for_engagement,
    list_visible_engagements,
    list_pending_download_requests,
    list_pending_storage_jobs,
    list_storage_jobs,
    list_storage_nodes,
    log_document_access,
    release_document_final,
    reopen_final_document,
    update_document_udin,
    user_can_manage_udin,
    user_can_view_udin_register,
    create_download_request_for_version,
    download_request_cache_path,
    fulfill_download_request_from_upload,
    save_uploaded_revision,
    user_can_delete_document,
    user_can_upload_to_engagement,
    user_can_view_engagement,
    version_absolute_path,
    PERMANENT_DOCUMENT_CATEGORIES,
    acknowledge_permanent_storage_job,
    create_permanent_download_request_for_version,
    fulfill_permanent_download_request_from_upload,
    get_latest_permanent_version,
    get_permanent_document,
    get_permanent_download_request_for_node,
    get_permanent_storage_job_for_node,
    get_permanent_version,
    list_pending_permanent_download_requests,
    list_pending_permanent_storage_jobs,
    list_permanent_documents_for_client,
    list_visible_clients_for_permanent_documents,
    permanent_download_request_cache_path,
    permanent_version_absolute_path,
    save_uploaded_permanent_revision,
    user_can_upload_client_documents,
    user_can_view_client_documents,
)
from app.modules.services.models import ClientServiceSubscription
from app.modules.core.tenancy.models import Branch, Tenant
from app.modules.core.tenancy.year_control import is_row_financial_year_locked
from app.modules.documents.agent_package import ERP_LOCAL_AGENT_VERSION, build_agent_env, build_preconfigured_agent_zip, build_agent_update_zip
from app.modules.documents.models import BranchStorageNode

from app.modules.services.task_documents import (
    get_task_with_subscription,
    get_task_document_requirement,
    list_documents_for_task,
    list_task_document_requirements,
    requirement_upload_status,
    save_uploaded_task_document,
    update_task_document_evidence_review,
)
from app.modules.services.execution import recalculate_task_aqmm_status

router = APIRouter(prefix="/documents", tags=["documents-ui"])
logger = logging.getLogger("audit_storage_agent.documents_ui")


def _normalise_document_text(value) -> str:
    return re.sub(r"[^a-z0-9]+", " ", str(value or "").lower()).strip()


def _permanent_onboarding_requirements(client) -> list[dict[str, object]]:
    """Return the client-master KYC checklist without creating a second document system."""
    constitution = _normalise_document_text(getattr(client, "client_type", ""))
    requirements: list[dict[str, object]] = [
        {
            "code": "pan",
            "label": "PAN / Income-tax identity proof",
            "categories": {"income tax", "kyc"},
            "keywords": {"pan", "income tax"},
        },
        {
            "code": "kyc",
            "label": "Authorised person / proprietor KYC",
            "categories": {"kyc"},
            "keywords": {"kyc", "aadhaar", "aadhar", "identity", "authorised signatory", "authorized signatory"},
        },
        {
            "code": "bank",
            "label": "Bank account proof",
            "categories": {"bank"},
            "keywords": {"bank", "cancelled cheque", "canceled cheque"},
        },
    ]

    if getattr(client, "gstin", None):
        requirements.append({
            "code": "gst",
            "label": "GST registration certificate",
            "categories": {"gst registration"},
            "keywords": {"gst", "gstin"},
        })

    if "partnership" in constitution:
        requirements.append({
            "code": "constitution",
            "label": "Partnership deed",
            "categories": {"agreements", "company registration"},
            "keywords": {"partnership deed", "deed"},
        })
    elif constitution == "llp" or "limited liability partnership" in constitution:
        requirements.extend([
            {
                "code": "registration",
                "label": "LLP incorporation / registration certificate",
                "categories": {"company registration", "roc master data"},
                "keywords": {"llp incorporation", "certificate of incorporation", "incorporation"},
            },
            {
                "code": "constitution",
                "label": "LLP agreement",
                "categories": {"agreements"},
                "keywords": {"llp agreement", "agreement"},
            },
        ])
    elif "limited company" in constitution or constitution in {"private limited company", "public limited company"}:
        requirements.extend([
            {
                "code": "registration",
                "label": "Certificate of Incorporation",
                "categories": {"company registration", "roc master data"},
                "keywords": {"certificate of incorporation", "incorporation"},
            },
            {
                "code": "constitution",
                "label": "MOA / AOA",
                "categories": {"company registration", "agreements", "roc master data"},
                "keywords": {"moa", "aoa", "memorandum", "articles of association"},
            },
        ])
    elif "trust" in constitution:
        requirements.append({
            "code": "constitution",
            "label": "Trust deed / registration document",
            "categories": {"agreements", "company registration", "licenses"},
            "keywords": {"trust deed", "trust registration"},
        })
    elif "society" in constitution:
        requirements.extend([
            {
                "code": "registration",
                "label": "Society registration certificate",
                "categories": {"company registration", "licenses"},
                "keywords": {"society registration", "registration certificate"},
            },
            {
                "code": "constitution",
                "label": "Bye-laws / governing document",
                "categories": {"agreements"},
                "keywords": {"bye laws", "byelaws", "governing document"},
            },
        ])
    elif "proprietorship" in constitution:
        requirements.append({
            "code": "constitution",
            "label": "Business constitution / registration proof",
            "categories": {"licenses", "company registration", "gst registration"},
            "keywords": {"proprietorship", "shop act", "udyam", "business registration"},
        })
    elif constitution in {"aop", "huf"}:
        requirements.append({
            "code": "constitution",
            "label": "Constitution / formation document",
            "categories": {"agreements", "kyc", "company registration"},
            "keywords": {"formation", "constitution", "deed", "huf"},
        })

    return requirements


def _build_permanent_onboarding_checklist(client, documents) -> tuple[list[dict[str, object]], bool]:
    haystacks = []
    for document in documents or []:
        haystacks.append({
            "category": _normalise_document_text(getattr(document, "category", "")),
            "text": _normalise_document_text(" ".join([
                str(getattr(document, "title", "") or ""),
                str(getattr(document, "description", "") or ""),
                str(getattr(document, "document_code", "") or ""),
            ])),
        })

    checklist = []
    for requirement in _permanent_onboarding_requirements(client):
        category_tokens = {_normalise_document_text(v) for v in requirement["categories"]}
        keyword_tokens = {_normalise_document_text(v) for v in requirement["keywords"]}
        complete = any(
            item["category"] in category_tokens
            or any(keyword and keyword in item["text"] for keyword in keyword_tokens)
            for item in haystacks
        )
        checklist.append({**requirement, "complete": complete})
    return checklist, all(item["complete"] for item in checklist)


def _base_ctx(request: Request, user, db, **ctx):
    base = {
        "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),
        "document_types": DOCUMENT_TYPES,
        "permanent_document_categories": PERMANENT_DOCUMENT_CATEGORIES,
    }
    base.update(ctx)
    return base


def _render(request: Request, template: str, db, user, **ctx):
    return templates.TemplateResponse(template, _base_ctx(request, user, db, **ctx))


def _redirect_denied():
    from app.core.http_responses import ui_access_denied
    return ui_access_denied()


def _require_user(request: Request, db, permission_code: str):
    user = get_current_user(request, db=db)
    if not user:
        return None, RedirectResponse(url="/login", status_code=303)
    try:
        require_permission(db, user, permission_code)
    except Exception:
        return user, _redirect_denied()
    return user, None




def _safe_documents_return_url(value: str | None, fallback: str) -> str:
    candidate = (value or "").strip()
    if candidate.startswith("/") and not candidate.startswith("//"):
        return candidate
    return fallback


def _active_financial_year(request: Request) -> str | None:
    value = request.session.get("active_financial_year") or getattr(request.state, "year_code", None)
    value = (value or "").strip()
    return value or None


def _parse_optional_date(value: str | None) -> date | None:
    value = (value or "").strip()
    if not value:
        return None
    return date.fromisoformat(value)


def _bool_from_form(value: str | None) -> bool:
    return (value or "").strip().lower() in {"1", "true", "yes", "on", "required"}


def _load_engagement(db, engagement_id: int):
    return db.execute(
        select(ClientServiceSubscription)
        .options(
            joinedload(ClientServiceSubscription.client),
            joinedload(ClientServiceSubscription.catalogue),
            joinedload(ClientServiceSubscription.assigned_partner),
            joinedload(ClientServiceSubscription.assigned_staff),
        )
        .where(ClientServiceSubscription.id == engagement_id)
    ).unique().scalar_one_or_none()


@router.get("")
def documents_home(request: Request, q: str = "", financial_year: str = ""):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        selected_financial_year = (financial_year or _active_financial_year(request) or "").strip()
        engagements = list_visible_engagements(db, user, scope, q=q, financial_year=selected_financial_year or None, limit=50)
        return _render(request, "modules/documents/templates/documents/index.html", db, user, title="Engagement Documents", q=q, financial_year=selected_financial_year, engagements=engagements)
    finally:
        db.close()


@router.get("/engagements/{engagement_id}")
def engagement_documents(request: Request, engagement_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        engagement = _load_engagement(db, engagement_id)
        active_fy = _active_financial_year(request)
        if active_fy and engagement and engagement.financial_year != active_fy:
            return RedirectResponse(url=f"/documents?financial_year={active_fy}", status_code=303)
        if not engagement or not user_can_view_engagement(db, user, engagement, scope):
            return _redirect_denied()
        documents = list_documents_for_engagement(db, engagement_id)
        download_requests = list_recent_download_requests_for_engagement(db, engagement_id, user_id=user.id, limit=10)
        return _render(request, "modules/documents/templates/documents/engagement_documents.html", db, user, title="Engagement Documents", engagement=engagement, documents=documents, download_requests=download_requests, can_upload=user_can_upload_to_engagement(db, user, engagement, scope))
    finally:
        db.close()


@router.post("/engagements/{engagement_id}/upload")
def upload_engagement_document(
    request: Request,
    engagement_id: int,
    title: str = Form(""),
    document_type: str = Form("GENERAL"),
    description: str | None = Form(None),
    remarks: str | None = Form(None),
    existing_document_id: str | None = Form(None),
    udin_required: str | None = Form(None),
    evidence_type: str | None = Form(None),
    evidence_description: str | None = Form(None),
    file: UploadFile = File(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        engagement = _load_engagement(db, engagement_id)
        active_fy = _active_financial_year(request)
        if active_fy and engagement and engagement.financial_year != active_fy:
            return RedirectResponse(url=f"/documents?financial_year={active_fy}", status_code=303)
        if not engagement or not user_can_upload_to_engagement(db, user, engagement, scope):
            return _redirect_denied()
        if is_row_financial_year_locked(db, engagement):
            return RedirectResponse(url=f"/documents/engagements/{engagement.id}?year_locked=1", status_code=303)
        if not file or not file.filename:
            return RedirectResponse(url=f"/documents/engagements/{engagement_id}?error=missing_file", status_code=303)
        try:
            doc = save_uploaded_revision(
                db,
                engagement=engagement,
                upload_file=file,
                title=title,
                document_type=document_type,
                description=description,
                remarks=remarks,
                user=user,
                existing_document_id=int(existing_document_id) if existing_document_id else None,
                udin_required=_bool_from_form(udin_required),
                evidence_type=evidence_type,
                evidence_description=evidence_description,
            )
            recalculate_task_aqmm_status(db, task)
            log_document_access(db, action="upload", result="success", user=user, request=request, document=doc)
            db.commit()
        except Exception as exc:
            db.rollback()
            logger.exception("Document upload failed for engagement_id=%s user_id=%s filename=%s", engagement_id, getattr(user, "id", None), getattr(file, "filename", None))
            try:
                log_document_access(db, action="upload", result="failed", user=user, request=request, message=str(exc)[:1000])
                db.commit()
            except Exception:
                db.rollback()
                logger.exception("Unable to write document upload failure audit log for engagement_id=%s", engagement_id)
            return RedirectResponse(url=f"/documents/engagements/{engagement_id}?error=upload_failed", status_code=303)
        return RedirectResponse(url=f"/documents/engagements/{engagement_id}?uploaded=1", status_code=303)
    finally:
        db.close()



@router.get("/tasks/{task_id}")
def task_documents(request: Request, task_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        task = get_task_with_subscription(db, task_id)
        engagement = task.subscription if task else None
        active_fy = _active_financial_year(request)
        if active_fy and engagement and engagement.financial_year != active_fy:
            return RedirectResponse(url=f"/documents?financial_year={active_fy}", status_code=303)
        if not task or not engagement or not user_can_view_engagement(db, user, engagement, scope):
            return _redirect_denied()
        requirements = []
        if task.firm_task_template_id:
            requirements = list_task_document_requirements(
                db,
                tenant_id=task.tenant_id,
                firm_task_template_id=task.firm_task_template_id,
            )
        documents = list_documents_for_task(db, task.id)
        requirement_status = requirement_upload_status(requirements, documents)
        return _render(
            request,
            "modules/documents/templates/documents/task_documents.html",
            db,
            user,
            title="Task Documents",
            task=task,
            engagement=engagement,
            requirements=requirements,
            requirement_status=requirement_status,
            documents=documents,
            can_upload=user_can_upload_to_engagement(db, user, engagement, scope),
            return_to=_safe_documents_return_url(request.query_params.get("return_to"), f"/documents/engagements/{engagement.id}"),
            back_url=_safe_documents_return_url(request.query_params.get("return_to"), f"/documents/engagements/{engagement.id}"),
        )
    finally:
        db.close()


@router.post("/tasks/{task_id}/upload")
def upload_task_document(
    request: Request,
    task_id: int,
    requirement_id: str | None = Form(None),
    title: str = Form(""),
    document_type: str = Form("GENERAL"),
    description: str | None = Form(None),
    remarks: str | None = Form(None),
    existing_document_id: str | None = Form(None),
    udin_required: str | None = Form(None),
    evidence_type: str | None = Form(None),
    evidence_description: str | None = Form(None),
    file: UploadFile = File(...),
    return_to: str = Form(""),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        return_target = _safe_documents_return_url(return_to, f"/documents/tasks/{task_id}")
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        task = get_task_with_subscription(db, task_id)
        engagement = task.subscription if task else None
        active_fy = _active_financial_year(request)
        if active_fy and engagement and engagement.financial_year != active_fy:
            return RedirectResponse(url=f"/documents?financial_year={active_fy}", status_code=303)
        if not task or not engagement or not user_can_upload_to_engagement(db, user, engagement, scope):
            return _redirect_denied()
        if is_row_financial_year_locked(db, engagement):
            return RedirectResponse(url=f"/documents/engagements/{engagement.id}?year_locked=1", status_code=303)
        if not file or not file.filename:
            return RedirectResponse(url=return_target if return_to else f"/documents/tasks/{task_id}?error=missing_file", status_code=303)
        requirement = None
        if requirement_id and str(requirement_id).strip():
            requirement = get_task_document_requirement(db, requirement_id=int(requirement_id), tenant_id=task.tenant_id)
            if not requirement or requirement.firm_task_template_id != task.firm_task_template_id:
                return RedirectResponse(url=return_target if return_to else f"/documents/tasks/{task_id}?error=invalid_requirement", status_code=303)
        try:
            doc = save_uploaded_task_document(
                db,
                task=task,
                requirement=requirement,
                upload_file=file,
                title=title,
                document_type=document_type,
                description=description,
                remarks=remarks,
                user=user,
                existing_document_id=int(existing_document_id) if existing_document_id else None,
                udin_required=_bool_from_form(udin_required),
                evidence_type=evidence_type,
                evidence_description=evidence_description,
            )
            recalculate_task_aqmm_status(db, task)
            log_document_access(db, action="task_upload", result="success", user=user, request=request, document=doc)
            db.commit()
        except Exception as exc:
            db.rollback()
            logger.exception("Task document upload failed for task_id=%s user_id=%s filename=%s", task_id, getattr(user, "id", None), getattr(file, "filename", None))
            try:
                log_document_access(db, action="task_upload", result="failed", user=user, request=request, message=str(exc)[:1000])
                db.commit()
            except Exception:
                db.rollback()
            return RedirectResponse(url=return_target if return_to else f"/documents/tasks/{task_id}?error=upload_failed", status_code=303)
        return RedirectResponse(url=return_target if return_to else f"/documents/tasks/{task_id}?uploaded=1", status_code=303)
    finally:
        db.close()

@router.post("/tasks/{task_id}/documents/{document_id}/evidence-review")
def update_task_document_evidence(
    request: Request,
    task_id: int,
    document_id: int,
    evidence_status: str = Form("uploaded"),
    evidence_type: str | None = Form(None),
    evidence_description: str | None = Form(None),
    evidence_review_note: str | None = Form(None),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        task = get_task_with_subscription(db, task_id)
        engagement = task.subscription if task else None
        if not task or not engagement or not user_can_upload_to_engagement(db, user, engagement, scope):
            return _redirect_denied()
        document = get_document(db, document_id)
        if not document or int(document.task_instance_id or 0) != int(task.id):
            return RedirectResponse(url=f"/documents/tasks/{task_id}?error=invalid_document", status_code=303)
        update_task_document_evidence_review(
            db,
            document=document,
            evidence_status=evidence_status,
            evidence_type=evidence_type,
            evidence_description=evidence_description,
            evidence_review_note=evidence_review_note,
            user=user,
        )
        recalculate_task_aqmm_status(db, task)
        log_document_access(db, action="task_evidence_review", result="success", user=user, request=request, document=document)
        db.commit()
        return RedirectResponse(url=f"/documents/tasks/{task_id}?evidence_reviewed=1", status_code=303)
    finally:
        db.close()


@router.get("/udin-register")
def udin_register(request: Request, financial_year: str = "", status: str = "", q: str = ""):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "udin.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        if not user_can_view_udin_register(db, user, scope):
            return _redirect_denied()
        selected_financial_year = (financial_year or _active_financial_year(request) or "").strip()
        rows = list_udin_records(
            db,
            user,
            scope,
            financial_year=selected_financial_year or None,
            status=status or None,
            q=q or None,
        )
        return _render(
            request,
            "modules/documents/templates/documents/udin_register.html",
            db,
            user,
            title="UDIN Register",
            rows=rows,
            financial_year=selected_financial_year,
            selected_status=status,
            q=q,
        )
    finally:
        db.close()


@router.post("/{document_id}/udin")
def update_document_udin_submit(
    request: Request,
    document_id: int,
    udin_required: str | None = Form(None),
    udin_number: str | None = Form(None),
    udin_date: str | None = Form(None),
    udin_document_date: str | None = Form(None),
    udin_document_type: str | None = Form(None),
    udin_financial_year: str | None = Form(None),
    udin_amount: str | None = Form(None),
    udin_notes: str | None = Form(None),
    redirect_to: str | None = Form(None),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "udin.manage")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        document = get_document(db, document_id)
        if not document or not user_can_manage_udin(db, user, document, scope):
            return _redirect_denied()
        if is_row_financial_year_locked(db, document.engagement):
            return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?year_locked=1", status_code=303)
        try:
            update_document_udin(
                db,
                document=document,
                udin_required=_bool_from_form(udin_required),
                udin_number=udin_number,
                udin_date=_parse_optional_date(udin_date),
                udin_document_date=_parse_optional_date(udin_document_date),
                udin_document_type=udin_document_type,
                udin_financial_year=udin_financial_year or document.financial_year,
                udin_amount=udin_amount,
                udin_notes=udin_notes,
                user=user,
            )
            log_document_access(db, action="udin_update", result="success", user=user, request=request, document=document)
            db.commit()
            target = (redirect_to or f"/documents/engagements/{document.engagement_id}").strip()
            return RedirectResponse(url=f"{target}{'&' if '?' in target else '?'}udin_saved=1", status_code=303)
        except Exception as exc:
            db.rollback()
            log_document_access(db, action="udin_update", result="failed", user=user, request=request, document=document, message=str(exc)[:1000])
            db.commit()
            return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?error=udin_failed", status_code=303)
    finally:
        db.close()


@router.post("/{document_id}/final-release")
def release_document_final_submit(
    request: Request,
    document_id: int,
    release_notes: str | None = Form(None),
    redirect_to: str | None = Form(None),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "udin.manage")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        document = get_document(db, document_id)
        if not document or not user_can_manage_udin(db, user, document, scope):
            return _redirect_denied()
        if is_row_financial_year_locked(db, document.engagement):
            return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?year_locked=1", status_code=303)
        try:
            release_document_final(db, document=document, release_notes=release_notes, user=user)
            log_document_access(db, action="final_release", result="success", user=user, request=request, document=document)
            db.commit()
            target = (redirect_to or f"/documents/engagements/{document.engagement_id}").strip()
            return RedirectResponse(url=f"{target}{'&' if '?' in target else '?'}released=1", status_code=303)
        except Exception as exc:
            db.rollback()
            log_document_access(db, action="final_release", result="failed", user=user, request=request, document=document, message=str(exc)[:1000])
            db.commit()
            target = (redirect_to or f"/documents/engagements/{document.engagement_id}").strip()
            return RedirectResponse(url=f"{target}{'&' if '?' in target else '?'}error=release_failed", status_code=303)
    finally:
        db.close()


@router.post("/{document_id}/reopen-final")
def reopen_final_document_submit(
    request: Request,
    document_id: int,
    release_notes: str | None = Form(None),
    redirect_to: str | None = Form(None),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "udin.manage")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        document = get_document(db, document_id)
        if not document or not user_can_manage_udin(db, user, document, scope):
            return _redirect_denied()
        reopen_final_document(db, document=document, release_notes=release_notes, user=user)
        log_document_access(db, action="final_reopen", result="success", user=user, request=request, document=document)
        db.commit()
        target = (redirect_to or f"/documents/engagements/{document.engagement_id}").strip()
        return RedirectResponse(url=f"{target}{'&' if '?' in target else '?'}reopened=1", status_code=303)
    finally:
        db.close()


def _stream_file(path: Path, download_name: str, content_type: str | None = None):
    def file_iterator():
        with path.open("rb") as fh:
            while True:
                chunk = fh.read(1024 * 1024)
                if not chunk:
                    break
                yield chunk

    quoted = quote(download_name)
    headers = {"Content-Disposition": f"attachment; filename*=UTF-8''{quoted}"}
    return StreamingResponse(file_iterator(), media_type=content_type or "application/octet-stream", headers=headers)


def _download_or_queue_from_local_node(request: Request, db, user, document, version, action: str):
    """Stream immediately if ERP/staged copy exists, otherwise queue DS5 local retrieval."""
    path = version_absolute_path(version)
    if path.exists():
        log_document_access(db, action=action, result="success", user=user, request=request, document=document, version=version, message="source=erp_stage")
        db.commit()
        return _stream_file(path, version.original_filename, version.content_type)

    download_request = create_download_request_for_version(db, version=version, user=user, request=request)
    if download_request:
        log_document_access(db, action=action, result="queued_local_retrieval", user=user, request=request, document=document, version=version, message=f"download_request_id={download_request.id}")
        db.commit()
        return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?download_queued={download_request.id}", status_code=303)

    log_document_access(db, action=action, result="missing_file", user=user, request=request, document=document, version=version, message="No ERP copy and no completed branch storage job found.")
    db.commit()
    return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?error=file_missing", status_code=303)


@router.get("/{document_id}/download")
def download_latest_document(request: Request, document_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.download")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        document = get_document(db, document_id)
        if not document:
            from app.core.http_responses import not_found_response
            return not_found_response(request, "Document not found")
        active_fy = _active_financial_year(request)
        if active_fy and document and getattr(document, "financial_year", None) != active_fy:
            return RedirectResponse(url=f"/documents?financial_year={active_fy}&error=wrong_year", status_code=303)
        if not document.engagement or not user_can_view_engagement(db, user, document.engagement, scope):
            return _redirect_denied()
        version = get_latest_version(document)
        if not version:
            return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?error=no_version", status_code=303)
        return _download_or_queue_from_local_node(request, db, user, document, version, action="download")
    finally:
        db.close()


@router.get("/versions/{version_id}/download")
def download_document_version(request: Request, version_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.download")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        version = get_version(db, version_id)
        if not version:
            from app.core.http_responses import not_found_response
            return not_found_response(request, "Document version not found")
        document = version.document
        active_fy = _active_financial_year(request)
        if active_fy and document and getattr(document, "financial_year", None) != active_fy:
            return RedirectResponse(url=f"/documents?financial_year={active_fy}&error=wrong_year", status_code=303)
        if not document or not document.engagement or not user_can_view_engagement(db, user, document.engagement, scope):
            return _redirect_denied()
        return _download_or_queue_from_local_node(request, db, user, document, version, action="download_version")
    finally:
        db.close()


@router.get("/download-requests")
def download_requests(request: Request, status: str = ""):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.download")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        requests = list_download_requests(db, tenant_id=_storage_tenant_filter(user, scope), branch_id=_storage_branch_filter(user, scope), status=status or None)
        return _render(request, "modules/documents/templates/documents/download_requests.html", db, user, title="Document Download Requests", requests=requests, status=status)
    finally:
        db.close()


@router.get("/download-requests/{request_id}/download")
def download_ready_request_file(request: Request, request_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.download")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        download_request = get_download_request(db, request_id)
        if not download_request or download_request.request_status != "ready":
            return RedirectResponse(url="/documents/download-requests?error=not_ready", status_code=303)
        document = download_request.document
        version = download_request.version
        active_fy = _active_financial_year(request)
        if active_fy and document and getattr(document, "financial_year", None) != active_fy:
            return RedirectResponse(url=f"/documents?financial_year={active_fy}&error=wrong_year", status_code=303)
        if not document or not version or not document.engagement or not user_can_view_engagement(db, user, document.engagement, scope):
            return _redirect_denied()
        path = download_request_cache_path(download_request)
        if not path or not path.exists():
            log_document_access(db, action="download_cached_request", result="missing_cache", user=user, request=request, document=document, version=version, message=f"download_request_id={download_request.id}")
            db.commit()
            return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?error=file_missing", status_code=303)
        log_document_access(db, action="download_cached_request", result="success", user=user, request=request, document=document, version=version, message=f"download_request_id={download_request.id}")
        db.commit()
        return _stream_file(path, version.original_filename, version.content_type)
    finally:
        db.close()





@router.get("/permanent")
def permanent_documents_home(request: Request, q: str = ""):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        clients = list_visible_clients_for_permanent_documents(db, user, scope, q=q, limit=200)
        return _render(request, "modules/documents/templates/documents/permanent_index.html", db, user, title="Permanent Client Documents", clients=clients, q=q)
    finally:
        db.close()


@router.get("/permanent/clients/{client_id}")
def permanent_client_documents(request: Request, client_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        from app.modules.clients.models import Client
        client = db.get(Client, client_id)
        if not client or not user_can_view_client_documents(db, user, client, scope):
            return _redirect_denied()
        documents = list_permanent_documents_for_client(db, client_id)
        onboarding_checklist, onboarding_complete = _build_permanent_onboarding_checklist(client, documents)
        return _render(
            request,
            "modules/documents/templates/documents/permanent_client_documents.html",
            db,
            user,
            title="Permanent Client Documents",
            client=client,
            documents=documents,
            can_upload=user_can_upload_client_documents(db, user, client, scope),
            onboarding_checklist=onboarding_checklist,
            onboarding_complete=onboarding_complete,
        )
    finally:
        db.close()


@router.post("/permanent/clients/{client_id}/upload")
def upload_permanent_client_document(
    request: Request,
    client_id: int,
    title: str = Form(""),
    category: str = Form("Other Permanent Documents"),
    description: str | None = Form(None),
    remarks: str | None = Form(None),
    existing_document_id: str | None = Form(None),
    file: UploadFile = File(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        from app.modules.clients.models import Client
        client = db.get(Client, client_id)
        if not client or not user_can_upload_client_documents(db, user, client, scope):
            return _redirect_denied()
        if not file or not file.filename:
            return RedirectResponse(url=f"/documents/permanent/clients/{client_id}?error=missing_file", status_code=303)
        try:
            doc = save_uploaded_permanent_revision(
                db,
                client=client,
                upload_file=file,
                title=title,
                category=category,
                description=description,
                remarks=remarks,
                user=user,
                existing_document_id=int(existing_document_id) if existing_document_id else None,
            )
            log_document_access(db, action="permanent_upload", result="success", user=user, request=request, message=f"permanent_document_id={doc.id}")
            db.commit()
        except Exception as exc:
            db.rollback()
            logger.exception("Permanent document upload failed for client_id=%s user_id=%s filename=%s", client_id, getattr(user, "id", None), getattr(file, "filename", None))
            try:
                log_document_access(db, action="permanent_upload", result="failed", user=user, request=request, message=str(exc)[:1000])
                db.commit()
            except Exception:
                db.rollback()
            return RedirectResponse(url=f"/documents/permanent/clients/{client_id}?error=upload_failed", status_code=303)
        return RedirectResponse(url=f"/documents/permanent/clients/{client_id}?uploaded=1", status_code=303)
    finally:
        db.close()


def _download_or_queue_permanent_from_local_node(request: Request, db, user, document, version):
    path = permanent_version_absolute_path(version)
    if path.exists():
        log_document_access(db, action="permanent_download", result="success", user=user, request=request, message=f"permanent_version_id={version.id};source=erp_stage")
        db.commit()
        return _stream_file(path, version.original_filename, version.content_type)
    download_request = create_permanent_download_request_for_version(db, version=version, user=user, request=request)
    if download_request:
        if getattr(download_request, "request_status", None) == "ready":
            cached_path = permanent_download_request_cache_path(download_request)
            if cached_path and cached_path.exists():
                log_document_access(db, action="permanent_download_cached_request", result="success", user=user, request=request, message=f"permanent_download_request_id={download_request.id}")
                db.commit()
                return _stream_file(cached_path, version.original_filename, version.content_type)
        log_document_access(db, action="permanent_download", result="queued_local_retrieval", user=user, request=request, message=f"permanent_download_request_id={download_request.id}")
        db.commit()
        return RedirectResponse(url=f"/documents/permanent/clients/{document.client_id}?download_queued={download_request.id}", status_code=303)
    log_document_access(db, action="permanent_download", result="missing_file", user=user, request=request, message=f"permanent_version_id={version.id}")
    db.commit()
    return RedirectResponse(url=f"/documents/permanent/clients/{document.client_id}?error=file_missing", status_code=303)


@router.get("/permanent/{document_id}/download")
def download_latest_permanent_document(request: Request, document_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.download")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        document = get_permanent_document(db, document_id)
        if not document:
            from app.core.http_responses import not_found_response
            return not_found_response(request, "Permanent document not found")
        if not document.client or not user_can_view_client_documents(db, user, document.client, scope):
            return _redirect_denied()
        version = get_latest_permanent_version(document)
        if not version:
            return RedirectResponse(url=f"/documents/permanent/clients/{document.client_id}?error=no_version", status_code=303)
        return _download_or_queue_permanent_from_local_node(request, db, user, document, version)
    finally:
        db.close()


@router.get("/permanent/versions/{version_id}/download")
def download_permanent_version(request: Request, version_id: int):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.download")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        version = get_permanent_version(db, version_id)
        if not version:
            from app.core.http_responses import not_found_response
            return not_found_response(request, "Permanent document version not found")
        document = version.document
        if not document or not document.client or not user_can_view_client_documents(db, user, document.client, scope):
            return _redirect_denied()
        return _download_or_queue_permanent_from_local_node(request, db, user, document, version)
    finally:
        db.close()


@router.post("/permanent/{document_id}/delete")
def delete_permanent_document(request: Request, document_id: int, csrf_token: str = Form(...)):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.delete")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        document = get_permanent_document(db, document_id)
        if not document or not document.client or not user_can_upload_client_documents(db, user, document.client, scope):
            return _redirect_denied()
        document.is_deleted = True
        document.status = "deleted"
        document.deleted_at_utc = datetime.now(timezone.utc)
        document.deleted_by_user_id = user.id
        document.updated_by_user_id = user.id
        log_document_access(db, action="permanent_delete", result="success", user=user, request=request, message=f"permanent_document_id={document.id}")
        db.commit()
        return RedirectResponse(url=f"/documents/permanent/clients/{document.client_id}?deleted=1", status_code=303)
    finally:
        db.close()

def _visible_branches(db, user, scope):
    stmt = select(Branch).where(Branch.is_active.is_(True)).order_by(Branch.tenant_id, Branch.name)
    forced_branch_id = _storage_branch_filter(user, scope)
    if forced_branch_id:
        stmt = stmt.where(Branch.id == int(forced_branch_id))
    elif not scope.is_system_admin:
        tenant_id = getattr(user, "tenant_id", None) or scope.tenant_id
        if tenant_id:
            stmt = stmt.where(Branch.tenant_id == int(tenant_id))
    return db.execute(stmt).scalars().all()


def _make_node_code(db, tenant_id: int, branch_id: int | None) -> str:
    tenant = db.get(Tenant, tenant_id)
    branch = db.get(Branch, branch_id) if branch_id else None
    tenant_code = re.sub(r"[^A-Za-z0-9]+", "", (getattr(tenant, "code", None) or f"AF{tenant_id}"))[:12].upper() or f"AF{tenant_id}"
    branch_code = re.sub(r"[^A-Za-z0-9]+", "", (getattr(branch, "code", None) or "FIRM"))[:12].upper() or "FIRM"
    base = f"{tenant_code}-{branch_code}-STORAGE"
    existing = {
        row[0]
        for row in db.execute(select(BranchStorageNode.node_code).where(BranchStorageNode.node_code.like(f"{base}%"))).all()
    }
    if base not in existing:
        return base
    for idx in range(2, 1000):
        candidate = f"{base}-{idx:03d}"
        if candidate not in existing:
            return candidate
    return f"{base}-{__import__('secrets').token_hex(3).upper()}"


def _default_node_name(db, tenant_id: int, branch_id: int | None) -> str:
    branch = db.get(Branch, branch_id) if branch_id else None
    if branch:
        return f"{branch.name} Local Storage"
    tenant = db.get(Tenant, tenant_id)
    return f"{getattr(tenant, 'name', 'Audit Firm')} Local Storage"


def _agent_download_filename(node_code: str, suffix: str) -> str:
    safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", node_code or "storage_node")
    return f"ERPLocalAgent_{safe}{suffix}"


def _find_existing_storage_node(db, tenant_id: int, branch_id: int | None) -> BranchStorageNode | None:
    """Return the canonical storage node for an audit-firm/branch pair.

    Important business rule:
        One Audit Firm + One Branch = One Storage Node Code forever.

    Earlier builds sometimes generated ARRR-01-STORAGE-002 / -003 for the
    same branch. That breaks the local .audit_storage_node.json identity file
    and can cause files to be stored through different agents. Therefore this
    function deliberately reuses the oldest row for the same tenant_id +
    branch_id, even if a newer duplicate is active. Package regeneration must
    rotate only the secret, not the node code.
    """
    stmt = select(BranchStorageNode).where(BranchStorageNode.tenant_id == int(tenant_id))
    if branch_id is None:
        stmt = stmt.where(BranchStorageNode.branch_id.is_(None))
    else:
        stmt = stmt.where(BranchStorageNode.branch_id == int(branch_id))
    stmt = stmt.order_by(BranchStorageNode.id.asc())
    return db.execute(stmt.limit(1)).scalar_one_or_none()


def _deactivate_duplicate_storage_nodes(db, node: BranchStorageNode) -> int:
    """Keep exactly one active storage node for one audit-firm/branch.

    Business rule: one branch can have only one permissible local storage agent.
    Any older duplicate nodes for the same tenant_id + branch_id are disabled so
    an old downloaded package/service cannot continue connecting to ERP.
    """
    stmt = select(BranchStorageNode).where(
        BranchStorageNode.tenant_id == int(node.tenant_id),
        BranchStorageNode.id != int(node.id),
        BranchStorageNode.is_active.is_(True),
    )
    if node.branch_id is None:
        stmt = stmt.where(BranchStorageNode.branch_id.is_(None))
    else:
        stmt = stmt.where(BranchStorageNode.branch_id == int(node.branch_id))
    duplicates = db.execute(stmt).scalars().all()
    for dup in duplicates:
        dup.is_active = False
        dup.status = "disabled_duplicate"
    return len(duplicates)


def _cleanup_visible_duplicate_storage_nodes(db, *, tenant_id: int | None, branch_id: int | None) -> None:
    """Keep the original node code as canonical and disable duplicates.

    For every audit-firm/branch pair, the oldest node row is treated as the
    permanent/canonical node. Newer duplicates are disabled, including rows that
    were previously generated with suffixes like -002. This matches the local
    storage identity file rule and prevents multiple branch agents for one
    branch.
    """
    stmt = select(BranchStorageNode)
    if tenant_id:
        stmt = stmt.where(BranchStorageNode.tenant_id == int(tenant_id))
    if branch_id:
        stmt = stmt.where(BranchStorageNode.branch_id == int(branch_id))
    nodes = db.execute(stmt.order_by(BranchStorageNode.tenant_id, BranchStorageNode.branch_id, BranchStorageNode.id.asc())).scalars().all()
    canonical_by_key: dict[tuple[int, int | None], BranchStorageNode] = {}
    for node in nodes:
        key = (int(node.tenant_id), int(node.branch_id) if node.branch_id is not None else None)
        canonical = canonical_by_key.get(key)
        if canonical is None:
            canonical_by_key[key] = node
            continue
        if node.is_active or node.status == "active":
            node.is_active = False
            node.status = "disabled_duplicate"



def _normalise_storage_root(value: str | None) -> str:
    cleaned = (value or "").strip().strip('"').strip("'")
    return cleaned or r"D:\AuditFirmStorage"


def _effective_storage_root(node: BranchStorageNode, requested_root: str | None = None, *, allow_change: bool = False) -> str:
    """Return the one permissible storage root for a branch node.

    Existing configured root always wins. This prevents repeated package
    generation or accidental reinstall from silently creating a second local
    storage folder for the same branch. A future dedicated root-change flow can
    pass allow_change=True after warning/confirmation.
    """
    requested = _normalise_storage_root(requested_root)
    existing = (node.storage_root_path or "").strip()
    if existing and not allow_change:
        return existing
    return requested or existing or r"D:\AuditFirmStorage"

def _update_node_secret_and_package(db, node: BranchStorageNode, *, storage_root_path: str | None, request: Request, include_admin_readme: bool = False, allow_storage_root_change: bool = False) -> Response:
    raw_secret = generate_storage_secret()
    node.secret_key_hash = hash_storage_secret(raw_secret)
    effective_root = _effective_storage_root(node, storage_root_path, allow_change=allow_storage_root_change)
    node.storage_root_path = effective_root
    node.is_active = True
    node.status = "active"
    _deactivate_duplicate_storage_nodes(db, node)
    db.flush()
    env_text = build_agent_env(
        erp_base_url=str(request.base_url).rstrip("/"),
        node_code=node.node_code,
        node_secret=raw_secret,
        storage_root=effective_root,
        tenant_id=node.tenant_id,
        branch_id=node.branch_id,
    )
    package = build_preconfigured_agent_zip(env_text=env_text, include_admin_readme=include_admin_readme)
    filename = _agent_download_filename(node.node_code, ".zip")
    return Response(
        package,
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )



def _is_partner_branch_storage_scope(scope) -> bool:
    """Partner/Branch Manager storage is limited to their own active/user branch."""
    return (scope.is_partner or scope.is_branch_manager) and not scope.is_firm_admin and not scope.is_system_admin


def _storage_branch_filter(user, scope) -> int | None:
    """Return branch id that should restrict storage-node screens for branch-managed roles."""
    if _is_partner_branch_storage_scope(scope):
        return scope.branch_id or getattr(user, "branch_id", None)
    return None


def _storage_tenant_filter(user, scope) -> int | None:
    if scope.is_system_admin:
        return None
    return getattr(user, "tenant_id", None) or scope.tenant_id


def _can_manage_branch_storage(scope) -> bool:
    # System Admin is monitoring/support-only for branch secrets/packages.
    return bool(scope.is_firm_admin or scope.is_partner or scope.is_branch_manager)


def _branch_name_map(db, branches=None):
    if branches is not None:
        return {b.id: b for b in branches}
    return {b.id: b for b in db.execute(select(Branch)).scalars().all()}


def _storage_scope_title(user, scope) -> str:
    if scope.is_system_admin:
        return "All audit firms â€” monitoring only"
    if scope.is_firm_admin:
        return "All branches of your audit firm"
    if _is_partner_branch_storage_scope(scope):
        branch_id = _storage_branch_filter(user, scope)
        return f"Your managed branch only{f' (Branch ID {branch_id})' if branch_id else ''}"
    return "Your permitted branch storage scope"


def _node_allowed_for_storage_scope(node: BranchStorageNode | None, user, scope) -> bool:
    if not node:
        return False
    tenant_filter = _storage_tenant_filter(user, scope)
    if tenant_filter and node.tenant_id != int(tenant_filter):
        return False
    branch_filter = _storage_branch_filter(user, scope)
    if branch_filter and node.branch_id != int(branch_filter):
        return False
    return True


def _selected_or_forced_branch_id(branch_id: str | None, user, scope) -> int | None:
    forced_branch_id = _storage_branch_filter(user, scope)
    if forced_branch_id:
        return int(forced_branch_id)
    return int(branch_id) if branch_id and str(branch_id).isdigit() else None

@router.get("/storage-nodes")
def storage_nodes(request: Request):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        tenant_filter = _storage_tenant_filter(user, scope)
        branch_filter = _storage_branch_filter(user, scope)
        _cleanup_visible_duplicate_storage_nodes(db, tenant_id=tenant_filter, branch_id=branch_filter)
        db.commit()
        nodes = list_storage_nodes(db, tenant_id=tenant_filter, branch_id=branch_filter)
        branches = _visible_branches(db, user, scope)
        jobs = list_storage_jobs(db, tenant_id=tenant_filter, branch_id=branch_filter, limit=10)
        requests = list_download_requests(db, tenant_id=tenant_filter, branch_id=branch_filter, limit=10)
        return _render(
            request,
            "modules/documents/templates/documents/storage_nodes.html",
            db,
            user,
            title="Branch Storage Nodes",
            nodes=nodes,
            branches=branches,
            branch_map=_branch_name_map(db, branches),
            recent_jobs=jobs,
            recent_download_requests=requests,
            generated_secret=None,
            storage_scope_title=_storage_scope_title(user, scope),
            erp_local_agent_version=ERP_LOCAL_AGENT_VERSION,
            can_manage_branch_storage=_can_manage_branch_storage(scope),
            forced_branch_id=branch_filter,
        )
    finally:
        db.close()


@router.get("/branch-storage-dashboard")
def partner_branch_storage_dashboard(request: Request):
    """Partner/Branch dashboard card for local storage setup and status.

    It intentionally reuses the same storage node data model and keeps the full
    Storage Nodes screen intact. Partner/Branch Manager users are automatically
    scoped to their own branch. Staff users receive a clear 403 instead of a
    generic dashboard page.
    """
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        if not _can_manage_branch_storage(scope):
            return _redirect_denied()
    finally:
        db.close()
    return storage_nodes(request)


@router.post("/storage-nodes")
def create_storage_node(
    request: Request,
    node_code: str = Form(...),
    node_name: str = Form(...),
    branch_id: str | None = Form(None),
    connector_url: str | None = Form(None),
    storage_root_path: str | None = Form(None),
    quota_limit_gb: str | None = Form(None),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        if not _can_manage_branch_storage(scope):
            return _redirect_denied()
        tenant_id = getattr(user, "tenant_id", None) or scope.tenant_id
        branch_id_int = _selected_or_forced_branch_id(branch_id, user, scope)
        if not tenant_id:
            return _redirect_denied()
        if branch_id_int:
            branch = db.get(Branch, branch_id_int)
            if not branch or (not scope.is_system_admin and branch.tenant_id != int(tenant_id)):
                return _redirect_denied()
            tenant_id = branch.tenant_id
        qgb = int(quota_limit_gb) if quota_limit_gb and quota_limit_gb.isdigit() else None
        existing = _find_existing_storage_node(db, int(tenant_id), branch_id_int)
        if existing:
            # Do not create ARRR-01-STORAGE-002 style duplicates. Reuse the
            # canonical branch node and update safe metadata only. Node code is
            # intentionally kept unchanged forever for the branch.
            existing.node_name = existing.node_name or node_name or _default_node_name(db, int(tenant_id), branch_id_int)
            existing.connector_url = (connector_url or "").strip() or existing.connector_url
            existing.storage_root_path = existing.storage_root_path or _normalise_storage_root(storage_root_path)
            existing.quota_limit_bytes = int(qgb) * 1024 * 1024 * 1024 if qgb else existing.quota_limit_bytes
            existing.is_active = True
            existing.status = "active"
            node = existing
        else:
            node, _secret = create_branch_storage_node(
                db,
                tenant_id=int(tenant_id),
                branch_id=branch_id_int,
                node_code=node_code,
                node_name=node_name,
                connector_url=connector_url,
                storage_root_path=storage_root_path,
                quota_limit_gb=qgb,
                user=user,
            )
        _deactivate_duplicate_storage_nodes(db, node)
        db.commit()
        return RedirectResponse(url="/documents/storage-nodes?created=1", status_code=303)
    except Exception:
        db.rollback()
        return RedirectResponse(url="/documents/storage-nodes?error=create_failed", status_code=303)
    finally:
        db.close()




@router.post("/storage-nodes/auto-setup")
def auto_setup_storage_node(
    request: Request,
    branch_id: str | None = Form(None),
    storage_root_path: str | None = Form(r"D:\AuditFirmStorage"),
    quota_limit_gb: str | None = Form(None),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        if not _can_manage_branch_storage(scope):
            return _redirect_denied()
        tenant_id = getattr(user, "tenant_id", None) or scope.tenant_id
        branch_id_int = _selected_or_forced_branch_id(branch_id, user, scope)
        if _is_partner_branch_storage_scope(scope) and not branch_id_int:
            return RedirectResponse(url="/documents/storage-nodes?error=branch_not_linked", status_code=303)
        if not tenant_id:
            return _redirect_denied()
        if branch_id_int:
            branch = db.get(Branch, branch_id_int)
            if not branch or (not scope.is_system_admin and branch.tenant_id != int(tenant_id)):
                return _redirect_denied()
            tenant_id = branch.tenant_id
        qgb = int(quota_limit_gb) if quota_limit_gb and str(quota_limit_gb).isdigit() else None
        existing = _find_existing_storage_node(db, int(tenant_id), branch_id_int)
        if existing:
            existing.node_name = existing.node_name or _default_node_name(db, int(tenant_id), branch_id_int)
            existing.storage_root_path = existing.storage_root_path or _normalise_storage_root(storage_root_path)
            existing.quota_limit_bytes = int(qgb) * 1024 * 1024 * 1024 if qgb else existing.quota_limit_bytes
            node = existing
        else:
            node_code = _make_node_code(db, int(tenant_id), branch_id_int)
            node_name = _default_node_name(db, int(tenant_id), branch_id_int)
            node, _unused_secret = create_branch_storage_node(
                db,
                tenant_id=int(tenant_id),
                branch_id=branch_id_int,
                node_code=node_code,
                node_name=node_name,
                connector_url=None,
                storage_root_path=_normalise_storage_root(storage_root_path),
                quota_limit_gb=qgb,
                user=user,
            )
        response = _update_node_secret_and_package(
            db,
            node,
            storage_root_path=_effective_storage_root(node, storage_root_path),
            request=request,
            include_admin_readme=bool(scope.is_system_admin),
        )
        db.commit()
        return response
    except Exception:
        db.rollback()
        return RedirectResponse(url="/documents/storage-nodes?error=auto_setup_failed", status_code=303)
    finally:
        db.close()


@router.post("/storage-nodes/download-env")
def download_storage_node_env(request: Request, csrf_token: str = Form(...)):
    """Do not expose standalone .env download to normal UI.

    Branch packages are generated server-side and include the .env internally.
    This keeps NODE_SECRET out of HTML source and avoids accidental sharing.
    """
    validate_csrf(request, csrf_token)
    return RedirectResponse(url="/documents/storage-nodes?error=env_download_disabled", status_code=303)


@router.post("/storage-nodes/download-agent-package")
def download_preconfigured_storage_agent(
    request: Request,
    node_code: str = Form(...),
    storage_root_path: str = Form(r"D:\AuditFirmStorage"),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        if not _can_manage_branch_storage(scope):
            return _redirect_denied()
        node = db.execute(select(BranchStorageNode).where(BranchStorageNode.node_code == node_code)).scalar_one_or_none()
        if not _node_allowed_for_storage_scope(node, user, scope):
            return _redirect_denied()
        canonical = _find_existing_storage_node(db, int(node.tenant_id), int(node.branch_id) if node.branch_id is not None else None)
        if canonical and canonical.id != node.id:
            node = canonical
        response = _update_node_secret_and_package(
            db,
            node,
            storage_root_path=_effective_storage_root(node, storage_root_path),
            request=request,
            include_admin_readme=bool(scope.is_system_admin),
        )
        db.commit()
        return response
    finally:
        db.close()


@router.post("/storage-nodes/{node_id}/download-agent-package")
def download_storage_node_agent_by_id(
    request: Request,
    node_id: int,
    storage_root_path: str = Form(r"D:\AuditFirmStorage"),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        if not _can_manage_branch_storage(scope):
            return _redirect_denied()
        node = db.get(BranchStorageNode, node_id)
        if not _node_allowed_for_storage_scope(node, user, scope):
            return _redirect_denied()
        canonical = _find_existing_storage_node(db, int(node.tenant_id), int(node.branch_id) if node.branch_id is not None else None)
        if canonical and canonical.id != node.id:
            node = canonical
        response = _update_node_secret_and_package(
            db,
            node,
            storage_root_path=_effective_storage_root(node, storage_root_path),
            request=request,
            include_admin_readme=bool(scope.is_system_admin),
        )
        db.commit()
        return response
    finally:
        db.close()

@router.post("/storage-nodes/{node_id}/toggle")
def toggle_storage_node(request: Request, node_id: int, csrf_token: str = Form(...)):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.upload")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        if not _can_manage_branch_storage(scope):
            return _redirect_denied()
        from app.modules.documents.models import BranchStorageNode
        node = db.get(BranchStorageNode, node_id)
        if not _node_allowed_for_storage_scope(node, user, scope):
            return _redirect_denied()
        canonical = _find_existing_storage_node(db, int(node.tenant_id), int(node.branch_id) if node.branch_id is not None else None)
        if canonical and canonical.id != node.id:
            # Duplicate rows must never be re-enabled. Enable/disable only the
            # canonical branch node to preserve one agent per branch.
            node.is_active = False
            node.status = "disabled_duplicate"
            node = canonical
        node.is_active = not bool(node.is_active)
        node.status = "active" if node.is_active else "disabled"
        if node.is_active:
            _deactivate_duplicate_storage_nodes(db, node)
        db.commit()
        return RedirectResponse(url="/documents/storage-nodes", status_code=303)
    finally:
        db.close()


@router.get("/storage-jobs")
def storage_jobs(request: Request, status: str = ""):
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.view")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        jobs = list_storage_jobs(db, tenant_id=_storage_tenant_filter(user, scope), branch_id=_storage_branch_filter(user, scope), status=status or None)
        return _render(request, "modules/documents/templates/documents/storage_jobs.html", db, user, title="Document Storage Jobs", jobs=jobs, status=status)
    finally:
        db.close()


def _agent_auth(db, request: Request, x_node_code: str | None, x_node_secret: str | None):
    node = authenticate_storage_node(db, x_node_code, x_node_secret, request=request)
    if not node:
        return None, JSONResponse({"ok": False, "error": "invalid_storage_node_credentials"}, status_code=401)
    return node, None


@router.get("/erp-local-agent/update-manifest")
def erp_local_agent_update_manifest(request: Request, x_node_code: str | None = Header(None), x_node_secret: str | None = Header(None)):
    import hashlib
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        package = build_agent_update_zip()
        db.commit()
        return {"ok": True, "agent_name": "ERP Local Agent", "latest_version": ERP_LOCAL_AGENT_VERSION, "sha256": hashlib.sha256(package).hexdigest(), "size_bytes": len(package), "download_url": "/documents/erp-local-agent/update-package", "channel": "stable"}
    finally:
        db.close()

@router.get("/erp-local-agent/update-package")
def erp_local_agent_update_package(request: Request, x_node_code: str | None = Header(None), x_node_secret: str | None = Header(None)):
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        package = build_agent_update_zip()
        db.commit()
        filename = f"ERP_Local_Agent_{ERP_LOCAL_AGENT_VERSION}.zip"
        return Response(package, media_type="application/zip", headers={"Content-Disposition": f'attachment; filename="{filename}"'})
    finally:
        db.close()
@router.post("/storage-agent/heartbeat")
async def storage_agent_heartbeat(request: Request, x_node_code: str | None = Header(None), x_node_secret: str | None = Header(None)):
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        cleanup = retry_verified_vps_cleanup_for_node(db, node, limit=50)
        db.commit()
        return {"ok": True, "node_code": node.node_code, "status": node.status, "vps_cleanup": cleanup}
    finally:
        db.close()



def _normal_storage_jobs_payload(jobs):
    return [
        {
            "job_id": job.id,
            "job_kind": "engagement",
            "version_id": job.version_id,
            "document_id": job.document_id,
            "target_relative_path": job.target_relative_path,
            "local_relative_path": job.target_relative_path,
            "file_size_bytes": job.file_size_bytes,
            "expected_hash_sha256": job.expected_hash_sha256,
            "download_url": f"/documents/storage-agent/jobs/{job.id}/download",
        }
        for job in jobs
    ]


def _permanent_storage_jobs_payload(jobs):
    return [
        {
            "job_id": f"P{job.id}",
            "job_kind": "permanent",
            "version_id": job.version_id,
            "document_id": job.document_id,
            "target_relative_path": job.target_relative_path,
            "local_relative_path": job.target_relative_path,
            "file_size_bytes": job.file_size_bytes,
            "expected_hash_sha256": job.expected_hash_sha256,
            "download_url": f"/documents/storage-agent/jobs/P{job.id}/download",
        }
        for job in jobs
    ]


def _normal_download_requests_payload(items):
    return [
        {
            "request_id": item.id,
            "request_kind": "engagement",
            "version_id": item.version_id,
            "document_id": item.document_id,
            "local_relative_path": item.local_relative_path,
            "expected_hash_sha256": item.expected_hash_sha256,
            "file_hash": item.expected_hash_sha256,
            "file_size_bytes": item.file_size_bytes,
            "upload_url": f"/documents/storage-agent/download-requests/{item.id}/upload",
        }
        for item in items
    ]


def _permanent_download_requests_payload(items):
    return [
        {
            "request_id": f"P{item.id}",
            "request_kind": "permanent",
            "version_id": item.version_id,
            "document_id": item.document_id,
            "local_relative_path": item.local_relative_path,
            "expected_hash_sha256": item.expected_hash_sha256,
            "file_hash": item.expected_hash_sha256,
            "file_size_bytes": item.file_size_bytes,
            "upload_url": f"/documents/storage-agent/download-requests/P{item.id}/upload",
        }
        for item in items
    ]


def _storage_agent_sync_payload(db, node):
    # Retry cleanup in tunnel mode as well as normal heartbeat mode.
    retry_verified_vps_cleanup_for_node(db, node, limit=25)
    jobs = _normal_storage_jobs_payload(list_pending_storage_jobs(db, node))
    jobs += _permanent_storage_jobs_payload(list_pending_permanent_storage_jobs(db, node))
    requests_ = _normal_download_requests_payload(list_pending_download_requests(db, node))
    requests_ += _permanent_download_requests_payload(list_pending_permanent_download_requests(db, node))
    return {"ok": True, "jobs": jobs, "download_requests": requests_, "requests": requests_}


@router.get("/storage-agent/jobs/pending")
def storage_agent_pending_jobs(request: Request, x_node_code: str | None = Header(None), x_node_secret: str | None = Header(None)):
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        payload = _storage_agent_sync_payload(db, node)
        db.commit()
        return {"ok": True, "jobs": payload["jobs"]}
    finally:
        db.close()


@router.get("/storage-agent/jobs/{job_id}/download")
def storage_agent_download_job(request: Request, job_id: str, x_node_code: str | None = Header(None), x_node_secret: str | None = Header(None)):
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        raw_job_id = str(job_id)
        is_permanent = raw_job_id.upper().startswith("P")
        actual_job_id = int(raw_job_id[1:] if is_permanent else raw_job_id)
        if is_permanent:
            job = get_permanent_storage_job_for_node(db, node, actual_job_id)
            path = permanent_version_absolute_path(job.version) if job and job.version else None
        else:
            job = get_storage_job_for_node(db, node, actual_job_id)
            path = version_absolute_path(job.version) if job and job.version else None
        if not job or job.status not in {"pending", "retry"}:
            return JSONResponse({"ok": False, "error": "job_not_available"}, status_code=404)
        if not path or not path.exists():
            job.status = "failed"
            job.last_error = "Staged file missing on ERP server."
            db.commit()
            return JSONResponse({"ok": False, "error": "staged_file_missing"}, status_code=404)
        job.status = "picked"
        job.picked_at_utc = datetime.now(timezone.utc)
        db.commit()
        return _stream_file(path, Path(job.target_relative_path).name, job.version.content_type if job.version else None)
    finally:
        db.close()


@router.post("/storage-agent/jobs/{job_id}/ack")
async def storage_agent_ack_job(request: Request, job_id: str, x_node_code: str | None = Header(None), x_node_secret: str | None = Header(None)):
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        payload = await request.json()
        raw_job_id = str(job_id)
        is_permanent = raw_job_id.upper().startswith("P")
        actual_job_id = int(raw_job_id[1:] if is_permanent else raw_job_id)
        acknowledged_hash = (payload.get("sha256") or payload.get("sha256_hash") or payload.get("file_hash") or "").strip()
        local_final_path = (payload.get("local_final_path") or payload.get("local_relative_path") or "").strip() or None
        if is_permanent:
            job = get_permanent_storage_job_for_node(db, node, actual_job_id)
            if not job:
                return JSONResponse({"ok": False, "error": "job_not_found"}, status_code=404)
            ok = acknowledge_permanent_storage_job(db, node=node, job=job, acknowledged_hash=acknowledged_hash, local_final_path=local_final_path, success=bool(payload.get("success", True)), error=payload.get("error"))
        else:
            job = get_storage_job_for_node(db, node, actual_job_id)
            if not job:
                return JSONResponse({"ok": False, "error": "job_not_found"}, status_code=404)
            ok = acknowledge_storage_job(db, node=node, job=job, acknowledged_hash=acknowledged_hash, local_final_path=local_final_path, success=bool(payload.get("success", True)), error=payload.get("error"))
        # First commit the verified local acknowledgement. Only after that
        # durable commit is the VPS staging file eligible for deletion.
        db.commit()
        cleanup = {"enabled": True, "deleted": False, "reason": "acknowledgement_failed"}
        if ok:
            if is_permanent:
                cleanup = cleanup_completed_permanent_storage_job_vps_stage(db, job)
            else:
                cleanup = cleanup_completed_storage_job_vps_stage(db, job)
            db.commit()
        return {"ok": ok, "job_status": job.status, "vps_cleanup": cleanup}
    finally:
        db.close()


@router.get("/storage-agent/download-requests/pending")
def storage_agent_pending_download_requests(request: Request, x_node_code: str | None = Header(None), x_node_secret: str | None = Header(None)):
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        payload = _storage_agent_sync_payload(db, node)
        db.commit()
        return {"ok": True, "download_requests": payload["download_requests"], "requests": payload["requests"]}
    finally:
        db.close()


@router.websocket("/storage-agent/tunnel")
async def storage_agent_tunnel(websocket: WebSocket):
    """Outbound local-agent tunnel.

    The branch PC opens this WebSocket connection to ERP. ERP never opens an
    inbound connection to the branch PC. File movement still uses the existing
    authenticated HTTP endpoints; this tunnel is the always-on control channel
    that pushes pending job/request notifications to the agent.
    """
    node_code = websocket.query_params.get("node_code") or websocket.headers.get("x-node-code")
    node_secret = websocket.query_params.get("node_secret") or websocket.headers.get("x-node-secret")
    await websocket.accept()

    db = CommonSessionLocal()
    try:
        node = authenticate_storage_node(db, node_code, node_secret, request=None)
        if not node:
            await websocket.send_json({"ok": False, "type": "error", "error": "invalid_storage_node_credentials"})
            await websocket.close(code=1008)
            return
        try:
            node.storage_mode = "tunnel"
        except Exception:
            pass
        db.commit()
        await websocket.send_json({"ok": True, "type": "connected", "node_code": node.node_code, "storage_mode": getattr(node, "storage_mode", "tunnel")})
    finally:
        db.close()

    last_push = 0.0
    while True:
        try:
            try:
                message = await asyncio.wait_for(websocket.receive_json(), timeout=5.0)
            except asyncio.TimeoutError:
                message = None

            db = CommonSessionLocal()
            try:
                node = authenticate_storage_node(db, node_code, node_secret, request=None)
                if not node:
                    await websocket.send_json({"ok": False, "type": "error", "error": "node_deactivated_or_invalid"})
                    await websocket.close(code=1008)
                    return
                if message and message.get("type") in {"heartbeat", "agent_status"}:
                    try:
                        node.storage_mode = "tunnel"
                    except Exception:
                        pass
                now = datetime.now(timezone.utc).timestamp()
                force = bool(message and message.get("type") in {"ready", "sync_now", "agent_status"})
                if force or now - last_push >= 5:
                    payload = _storage_agent_sync_payload(db, node)
                    payload.update({"type": "sync", "node_code": node.node_code})
                    await websocket.send_json(payload)
                    last_push = now
                db.commit()
            finally:
                db.close()
        except WebSocketDisconnect:
            break
        except Exception as exc:
            logger.exception("Storage agent tunnel failed for node=%s: %s", node_code, exc)
            try:
                await websocket.send_json({"ok": False, "type": "error", "error": "tunnel_server_error"})
            except Exception:
                pass
            await asyncio.sleep(5)


@router.post("/storage-agent/download-requests/{request_id}/upload")
def storage_agent_upload_download_request(
    request: Request,
    request_id: str,
    file: UploadFile = File(...),
    x_node_code: str | None = Header(None),
    x_node_secret: str | None = Header(None),
):
    db = CommonSessionLocal()
    try:
        node, error = _agent_auth(db, request, x_node_code, x_node_secret)
        if error:
            return error
        raw_request_id = str(request_id)
        is_permanent = raw_request_id.upper().startswith("P")
        actual_request_id = int(raw_request_id[1:] if is_permanent else raw_request_id)
        if is_permanent:
            download_request = get_permanent_download_request_for_node(db, node, actual_request_id)
            if not download_request or download_request.request_status not in {"pending", "picked", "retry"}:
                return JSONResponse({"ok": False, "error": "download_request_not_available"}, status_code=404)
            download_request.request_status = "picked"
            ok = fulfill_permanent_download_request_from_upload(db, node=node, download_request=download_request, upload_file=file)
        else:
            download_request = get_download_request_for_node(db, node, actual_request_id)
            if not download_request or download_request.request_status not in {"pending", "picked", "retry"}:
                return JSONResponse({"ok": False, "error": "download_request_not_available"}, status_code=404)
            download_request.request_status = "picked"
            ok = fulfill_download_request_from_upload(db, node=node, download_request=download_request, upload_file=file)
        db.commit()
        return {"ok": ok, "request_status": download_request.request_status}
    finally:
        db.close()


@router.post("/{document_id}/delete")
def delete_document(request: Request, document_id: int, csrf_token: str = Form(...)):
    validate_csrf(request, csrf_token)
    db = CommonSessionLocal()
    try:
        user, response = _require_user(request, db, "documents.delete")
        if response:
            return response
        scope = build_document_scope(request, db, user)
        document = get_document(db, document_id)
        if not document or not user_can_delete_document(db, user, document, scope):
            return _redirect_denied()
        document.is_deleted = True
        document.status = "deleted"
        document.deleted_at_utc = datetime.now(timezone.utc)
        document.deleted_by_user_id = user.id
        document.updated_by_user_id = user.id
        log_document_access(db, action="delete", result="success", user=user, request=request, document=document)
        db.commit()
        return RedirectResponse(url=f"/documents/engagements/{document.engagement_id}?deleted=1", status_code=303)
    finally:
        db.close()

