arrr-erp/app/modules/accounting/gst_reconciliation_ui.py

from __future__ import annotations

import calendar
import re
import tempfile
import uuid
from datetime import date, datetime, timedelta, timezone
from pathlib import Path
from urllib.parse import urlencode

import jwt
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.security.csrf import get_or_create_csrf_token, validate_csrf
from app.core.settings import get_settings
from app.core.templating import templates
from app.modules.accounting.agent_bridge import request_agent_command
from app.modules.accounting.ui import _find_visible_client, _require_partner, _visible_clients, _node_online
from app.modules.core.rbac.deps import get_user_permissions, get_user_roles
from app.modules.credential_vault.crypto import decrypt_value
from app.modules.credential_vault.models import CredentialVaultEntry
from app.modules.credential_vault.service import can_view_entry, log_access
from app.modules.documents.services import client_folder_parts, get_active_storage_node_for_branch, sanitize_segment
from app.modules.registrations.models import ClientRegistration, RegistrationType

router = APIRouter(prefix="/tools/accounting/gst-reconciliation", tags=["accounting-gst-reconciliation-ui"])
_TOKEN_PURPOSE = "gst_operator_browser_v1"
_TOKEN_MINUTES = 30
_UPLOAD_ROOT = Path(tempfile.gettempdir()) / "audit_firm_gst_operator_uploads"


def _fy_bounds(fy: str) -> tuple[date, date]:
    m = re.fullmatch(r"(\d{4})-(\d{2})", str(fy or "").strip())
    if not m:
        raise ValueError("Invalid financial year.")
    y = int(m.group(1))
    return date(y, 4, 1), date(y + 1, 3, 31)


def _periods_for_fy(fy: str) -> list[str]:
    start, _ = _fy_bounds(fy)
    periods=[]
    y=start.year; m=4
    for _ in range(12):
        periods.append(f"{m:02d}{y:04d}")
        m += 1
        if m == 13:
            m = 1; y += 1
    return periods


def _fy_for_period(period: str) -> str:
    digits = re.sub(r"\D", "", period or "")
    if len(digits) != 6:
        raise ValueError("Return period must be MMYYYY.")
    month, year = int(digits[:2]), int(digits[2:])
    if month < 1 or month > 12:
        raise ValueError("Invalid GST return month.")
    sy = year if month >= 4 else year - 1
    return f"{sy}-{str(sy+1)[-2:]}"


def _period_bounds(period: str) -> tuple[str, str]:
    digits = re.sub(r"\D", "", period or "")
    month, year = int(digits[:2]), int(digits[2:])
    last = calendar.monthrange(year, month)[1]
    return date(year, month, 1).isoformat(), date(year, month, last).isoformat()


def _gst_regs(db, tenant_id: int, client_id: int):
    return db.execute(
        select(ClientRegistration, RegistrationType)
        .join(RegistrationType, RegistrationType.id == ClientRegistration.registration_type_id)
        .where(ClientRegistration.tenant_id == tenant_id, ClientRegistration.client_id == client_id)
        .order_by(ClientRegistration.id.asc())
    ).all()


def _is_gstin(reg, typ) -> bool:
    code = str(getattr(typ, "code", "") or "").upper().strip()
    num = re.sub(r"\s+", "", str(getattr(reg, "registration_number", "") or "").upper())
    return code == "GSTIN" and len(num) == 15


def _norm_gstin(value: str | None) -> str:
    return re.sub(r"[^0-9A-Z]", "", str(value or "").upper())


def _credential_is_gst_portal(entry: CredentialVaultEntry) -> bool:
    category = str(entry.category or "").strip().lower()
    title = str(entry.title or "").strip().lower()
    portal_url = str(entry.portal_url or "").strip().lower()
    reference = str(entry.reference_number or "").strip().lower()
    if category == "gst_portal":
        return True
    haystack = " ".join((title, portal_url, reference))
    gst_hint = (
        "gst portal" in haystack or "gst login" in haystack
        or "services.gst.gov.in" in portal_url or "www.gst.gov.in" in portal_url
    )
    if not gst_hint:
        return False
    if category in {"eway_bill", "einvoice", "api_key"}:
        return "gst portal" in title or "gst login" in title
    return True


def _credential_matches_gstin(entry: CredentialVaultEntry, registration_id: int, gstin: str) -> bool:
    if int(entry.registration_id or 0) == int(registration_id):
        return True
    if _norm_gstin(entry.reference_number) == _norm_gstin(gstin):
        return True
    return entry.registration_id is None


def _vault_entries(db, user, request, tenant_id: int, client_id: int, registration_id: int | None = None, gstin: str = ""):
    q = select(CredentialVaultEntry).where(
        CredentialVaultEntry.tenant_id == tenant_id,
        CredentialVaultEntry.client_id == client_id,
        CredentialVaultEntry.status != "archived",
    )
    rows = db.execute(q.order_by(CredentialVaultEntry.title.asc())).scalars().all()
    branch_id = request.session.get("active_branch_id") or getattr(user, "branch_id", None)
    visible = [r for r in rows if can_view_entry(db, user, r, int(branch_id) if branch_id else None)]
    if not registration_id:
        return [r for r in visible if _credential_is_gst_portal(r)]
    eligible = [r for r in visible if _credential_is_gst_portal(r) and _credential_matches_gstin(r, int(registration_id), gstin)]
    def rank(entry: CredentialVaultEntry):
        exact_registration = int(entry.registration_id or 0) == int(registration_id)
        exact_gstin = _norm_gstin(entry.reference_number) == _norm_gstin(gstin)
        client_level = entry.registration_id is None
        return (0 if exact_registration else 1 if exact_gstin else 2 if client_level else 3, str(entry.title or "").lower(), int(entry.id or 0))
    return sorted(eligible, key=rank)


def _storage_payload(client, fy: str, gstin: str) -> tuple[str, str]:
    fy_folder = sanitize_segment(f"FY{fy}", "FY")
    letter, client_folder = client_folder_parts(client, int(client.id))
    root = Path(fy_folder) / "Clients" / letter / client_folder
    return (root / "Accounting").as_posix(), (root / "GST" / sanitize_segment(f"GSTIN_{gstin}", "GSTIN")).as_posix()


def _redirect(client_id: int, **params):
    data = {"client_id": client_id, **{k: v for k, v in params.items() if v not in (None, "")}}
    return RedirectResponse("/tools/accounting/gst-reconciliation?" + urlencode(data), status_code=303)


def _encode_operator_token(payload: dict) -> str:
    now = datetime.now(timezone.utc)
    body = {**payload, "purpose": _TOKEN_PURPOSE, "iat": now, "exp": now + timedelta(minutes=_TOKEN_MINUTES)}
    return jwt.encode(body, get_settings().SECRET_KEY, algorithm="HS256")


def _decode_operator_token(token: str) -> dict:
    data = jwt.decode(token, get_settings().SECRET_KEY, algorithms=["HS256"])
    if data.get("purpose") != _TOKEN_PURPOSE:
        raise ValueError("Invalid GST operator token.")
    return data


def _selected_periods(download_mode: str, financial_year: str, period: str) -> list[str]:
    if download_mode == "full_fy":
        return _periods_for_fy(financial_year)
    digits = re.sub(r"\D", "", period or "")
    if len(digits) != 6:
        raise ValueError("Enter a valid MMYYYY period for Single Month mode.")
    return [digits]


@router.get("")
def page(request: Request, client_id: int | None = None, registration_id: int | None = None, period: str = "", financial_year: str = "", download_mode: str = "single", message: str = "", error: str = ""):
    db = CommonSessionLocal()
    try:
        user, response = _require_partner(request, db, "accounting.learning.view")
        if response:
            return response
        clients, scope = _visible_clients(db, request, user)
        selected = next((c for c in clients if client_id and int(c.id) == int(client_id)), None)
        registrations=[]; selected_reg=None; credentials=[]; node=None
        if not financial_year:
            financial_year = str(request.session.get("active_financial_year") or request.session.get("active_fy") or "2025-26")
        if download_mode not in {"single", "full_fy"}:
            download_mode = "single"
        if selected:
            registrations=[(r,t) for r,t in _gst_regs(db,scope.tenant_id,selected.id) if _is_gstin(r,t)]
            selected_reg=next((r for r,t in registrations if registration_id and int(r.id)==int(registration_id)),None)
            if not selected_reg and registrations:
                selected_reg=registrations[0][0]
            selected_gstin = _norm_gstin(getattr(selected_reg, "registration_number", "")) if selected_reg else ""
            credentials=_vault_entries(db,user,request,scope.tenant_id,selected.id,int(selected_reg.id) if selected_reg else None, selected_gstin)
            node=get_active_storage_node_for_branch(db,scope.tenant_id,scope.branch_id)
        operator_job = request.session.pop("gst_operator_job", None)
        return templates.TemplateResponse("modules/accounting/templates/accounting/gst_reconciliation.html",{
            "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),
            "clients":clients,"selected_client":selected,"registrations":registrations,"selected_registration":selected_reg,"credentials":credentials,"node":node,"node_online":_node_online(node) if node else False,
            "period":period,"financial_year":financial_year,"download_mode":download_mode,"operator_job":operator_job,"message":message,"error":error,"title":"GST Return Reconciliation",
        })
    finally:
        db.close()


@router.post("/download/start")
def start_download(
    request: Request,
    client_id: int=Form(...), registration_id: int=Form(...), credential_id: int=Form(...),
    period: str=Form(""), financial_year: str=Form(...), download_mode: str=Form("single"),
    gstr1: str=Form(""), gstr2b: str=Form(""), gstr3b: str=Form(""), gstr2a: str=Form(""),
    csrf_token: str=Form(...),
):
    validate_csrf(request,csrf_token)
    db=CommonSessionLocal()
    try:
        user,response=_require_partner(request,db,"accounting.learning.manage")
        if response: return response
        client,_,scope=_find_visible_client(db,request,user,client_id)
        if not client: return _redirect(client_id,error="Client is not available in your scope.")
        pair=next(((r,t) for r,t in _gst_regs(db,scope.tenant_id,client.id) if int(r.id)==registration_id and _is_gstin(r,t)),None)
        if not pair: return _redirect(client_id,error="Select a valid GSTIN registration.")
        reg,_=pair; gstin=_norm_gstin(reg.registration_number)
        cred=db.get(CredentialVaultEntry,credential_id)
        credential_ok = bool(cred and int(cred.tenant_id or 0)==int(scope.tenant_id) and int(cred.client_id or 0)==int(client.id) and str(cred.status or "").lower()!="archived" and can_view_entry(db,user,cred,scope.branch_id) and _credential_is_gst_portal(cred) and _credential_matches_gstin(cred,int(reg.id),gstin))
        if not credential_ok:
            return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,download_mode=download_mode,error="Selected GST Portal credential is not available for this client/GSTIN.")
        node=get_active_storage_node_for_branch(db,scope.tenant_id,scope.branch_id)
        if not node or not _node_online(node):
            return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,download_mode=download_mode,error="Local Storage Agent is offline.")
        periods=_selected_periods(download_mode,financial_year,period)
        return_types=[name for name,flag in (("GSTR1",gstr1),("GSTR2B",gstr2b),("GSTR3B",gstr3b),("GSTR2A",gstr2a)) if flag]
        if not return_types:
            return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,download_mode=download_mode,error="Select at least one GST return to download.")
        # Ensure FY is consistent in single-month mode.
        if download_mode == "single":
            financial_year=_fy_for_period(periods[0])
        accounting_dir,gst_dir=_storage_payload(client,financial_year,gstin)
        jti=uuid.uuid4().hex
        token=_encode_operator_token({
            "jti":jti,"user_id":int(user.id),"tenant_id":int(scope.tenant_id),"branch_id":int(scope.branch_id),
            "node_code":str(node.node_code),"client_id":int(client.id),"client_name":str(client.client_name or ""),
            "registration_id":int(reg.id),"credential_id":int(cred.id),"gstin":gstin,"financial_year":financial_year,
            "periods":periods,"return_types":return_types,"gst_relative_dir":gst_dir,"accounting_relative_dir":accounting_dir,
        })
        log_access(db,request,user,cred,"use_for_gst_download",reason=f"GST returns {financial_year}: {','.join(return_types)}",fields="username,secret",success=True)
        db.commit()
        request.session["gst_operator_job"]={"token":token,"job_id":jti,"periods":periods,"return_types":return_types}
        return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,download_mode=download_mode,message="GST download prepared. The ERP will now ask the Local Agent on this computer to open the visible GST browser; CAPTCHA/OTP will appear here. Completed files will be transferred to the configured client local storage.")
    except Exception as exc:
        db.rollback(); return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,download_mode=download_mode,error=str(exc))
    finally:
        db.close()


@router.post("/operator/redeem")
async def operator_redeem(request: Request):
    body=await request.json(); token=str(body.get("token") or "")
    try:
        data=_decode_operator_token(token)
        db=CommonSessionLocal()
        try:
            cred=db.get(CredentialVaultEntry,int(data["credential_id"]))
            if not cred or int(cred.tenant_id or 0)!=int(data["tenant_id"]) or int(cred.client_id or 0)!=int(data["client_id"]) or str(cred.status or "").lower()=="archived":
                raise ValueError("GST credential is no longer available.")
            username=decrypt_value(cred.tenant_id,cred.username_encrypted) or ""
            password=decrypt_value(cred.tenant_id,cred.secret_encrypted) or ""
            if not username or not password:
                raise ValueError("GST username/password is missing in Credential Vault.")
        finally:
            db.close()
        return JSONResponse({"ok":True,"payload":{**{k:v for k,v in data.items() if k not in {"iat","exp","purpose"}},"username":username,"password":password,"upload_url":str(get_settings().ERP_PUBLIC_BASE_URL).rstrip("/")+"/tools/accounting/gst-reconciliation/operator/upload","login_timeout_seconds":900}})
    except Exception as exc:
        return JSONResponse({"ok":False,"error":str(exc)},status_code=400)


@router.post("/operator/upload")
async def operator_upload(token: str=Form(...), package: UploadFile=File(...)):
    try:
        data=_decode_operator_token(token); jti=str(data.get("jti") or "")
        if not jti:
            raise ValueError("GST operator job id is missing.")
        _UPLOAD_ROOT.mkdir(parents=True,exist_ok=True)
        package_path=_UPLOAD_ROOT/f"{jti}.zip"
        total=0
        with package_path.open("wb") as out:
            while True:
                chunk=await package.read(1024*1024)
                if not chunk: break
                total += len(chunk)
                if total > 250*1024*1024:
                    raise ValueError("GST download package exceeds the 250 MB safety limit.")
                out.write(chunk)
        package_url=str(get_settings().ERP_PUBLIC_BASE_URL).rstrip("/")+f"/tools/accounting/gst-reconciliation/operator/package/{jti}?token={token}"
        result=request_agent_command(str(data["node_code"]),"gst_return_package_store",{
            "client_id":int(data["client_id"]),"gstin":str(data["gstin"]),"financial_year":str(data["financial_year"]),
            "gst_relative_dir":str(data["gst_relative_dir"]),"package_url":package_url,"periods":data.get("periods") or [],"return_types":data.get("return_types") or [],
        },timeout_seconds=120)
        if not result.get("ok"):
            raise RuntimeError(result.get("error") or "Local Storage Agent could not store the GST return package.")
        try: package_path.unlink(missing_ok=True)
        except Exception: pass
        return JSONResponse({"ok":True,"stored":result.get("result") or {}})
    except Exception as exc:
        return JSONResponse({"ok":False,"error":str(exc)},status_code=400)


@router.get("/operator/package/{job_id}")
def operator_package(job_id: str, token: str):
    try:
        data=_decode_operator_token(token)
        if str(data.get("jti") or "") != str(job_id):
            raise ValueError("GST package token does not match the requested job.")
        path=_UPLOAD_ROOT/f"{job_id}.zip"
        if not path.is_file():
            return JSONResponse({"ok":False,"error":"GST package is no longer available."},status_code=404)
        return FileResponse(path,media_type="application/zip",filename=f"gst_returns_{job_id}.zip")
    except Exception as exc:
        return JSONResponse({"ok":False,"error":str(exc)},status_code=400)


@router.post("/analyze")
def analyze(request: Request, client_id: int=Form(...), registration_id: int=Form(...), period: str=Form(""), financial_year: str=Form(...), analyze_mode: str=Form("single"), csrf_token: str=Form(...)):
    validate_csrf(request,csrf_token)
    db=CommonSessionLocal()
    try:
        user,response=_require_partner(request,db,"accounting.learning.manage")
        if response: return response
        client,_,scope=_find_visible_client(db,request,user,client_id)
        if not client: return _redirect(client_id,error="Client is not available in your scope.")
        pair=next(((r,t) for r,t in _gst_regs(db,scope.tenant_id,client.id) if int(r.id)==registration_id and _is_gstin(r,t)),None)
        if not pair: return _redirect(client_id,error="GSTIN registration was not found.")
        reg,_=pair; gstin=_norm_gstin(reg.registration_number)
        periods=_periods_for_fy(financial_year) if analyze_mode=="full_fy" else [re.sub(r"\D","",period or "")]
        if any(len(p)!=6 for p in periods): raise ValueError("Enter a valid MMYYYY period.")
        accounting_dir,gst_dir=_storage_payload(client,financial_year,gstin)
        node=get_active_storage_node_for_branch(db,scope.tenant_id,scope.branch_id)
        if not node or not _node_online(node): return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,error="Local Storage Agent is offline.")
        results=[]
        for p in periods:
            date_from,date_to=_period_bounds(p)
            res=request_agent_command(node.node_code,"gst_reconciliation_analyze",{"client_id":client.id,"gstin":gstin,"financial_year":financial_year,"period":p,"gst_relative_dir":gst_dir,"accounting_relative_dir":accounting_dir,"date_from":date_from,"date_to":date_to},timeout_seconds=30)
            if res.get("ok"):
                results.append((res.get("result") or {}).get("analysis") or {})
        if not results: raise RuntimeError("No stored GST periods could be reconciled.")
        if analyze_mode=="full_fy":
            def sum_counts(section):
                out={}
                for r in results:
                    for k,v in (((r.get(section) or {}).get("counts") or {}).items()): out[k]=out.get(k,0)+int(v or 0)
                return out
            taxes={k:{"gstr2b":0.0,"gstr3b":0.0,"difference":0.0} for k in ("igst","cgst","sgst","cess")}
            for r in results:
                for k,row in (r.get("itc_reconciliation") or {}).items():
                    if k in taxes:
                        for f in taxes[k]: taxes[k][f]=round(taxes[k][f]+float(row.get(f) or 0),2)
            analysis={"period":financial_year,"sales_reconciliation":{"counts":sum_counts("sales_reconciliation")},"purchase_reconciliation":{"counts":sum_counts("purchase_reconciliation")},"itc_reconciliation":taxes}
        else:
            analysis=results[0]
        request.session["gst_reconciliation_result"]=analysis
        return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,download_mode=("full_fy" if analyze_mode=="full_fy" else "single"),message="GST Purchase, Sales and ITC reconciliation completed from stored local return data and Accounting Mirror.")
    except Exception as exc:
        return _redirect(client_id,registration_id=registration_id,period=period,financial_year=financial_year,error=str(exc))
    finally:
        db.close()