Ten31-Portal/backend/ten31portal/routers/capital_import_router.py

"""Import fund members and their capital figures from a fund-administrator workbook.

Primary path: an eNAV workbook's "ALLOC SI" tab (one row per investor with INVESTOR ID,
INVESTOR NAME, COMMITTED CAPITAL, CONTRIBUTIONS, (DISTRIBUTIONS), ENDING BALANCE).
Fallback: a subaccounts sheet (investor names across columns, per-vehicle value rows).

Encrypted workbooks are decrypted with the open password. Nothing is written on preview.
Commit matches existing members (by fund-admin investor ID, else name), creates new ones
(on the shared default password unless one is given), grants entity access, and loads each
member's capital-account statement (commitment, contributions, distributions, current value).
"""

import io
import re
from datetime import date, datetime

import openpyxl
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from sqlmodel import Session, select

from ten31portal import config, storage
from ten31portal.audit import record_audit
from ten31portal.auth import hash_password, require_internal_admin
from ten31portal.database import get_session
from ten31portal.models import (
    CapitalAccountStatement, Entity, EntityAccess, User, UserRole,
)
from ten31portal.routers.import_router import _open_workbook, _enav_as_of
from ten31portal.schemas import (
    BatchCapitalFileResult, BatchCapitalImportResult,
    CapitalImportCommit, CapitalImportPreview, ImportInvestorPreview, ImportValueRow,
)

router = APIRouter(prefix="/api/import/capital-accounts", tags=["import"])

MAX_ROWS = 400
MAX_COLS = 90


def reset_entity_partners(entity_id: int, session: Session) -> dict[str, int]:
    """Remove every partner from one fund: delete its capital-account statements and the
    investors' access grants to it. The investor *accounts* are kept — they usually also
    belong to other funds — only their membership of THIS entity is cleared. Use to undo a
    wrong members/ALLOC-SI import (e.g. Fund II's roster loaded into Fund III). Holdings/NAV
    are untouched (see reset_entity_holdings for those). Caller commits.
    """
    statements = 0
    for s in session.exec(
        select(CapitalAccountStatement).where(
            CapitalAccountStatement.entity_id == entity_id
        )
    ).all():
        session.delete(s)
        statements += 1

    # Only drop investor memberships; a fund_administrator's access is not a "partner".
    investor_ids = {
        u.id for u in session.exec(
            select(User).where(User.role == UserRole.investor)
        ).all()
    }
    access = 0
    for a in session.exec(
        select(EntityAccess).where(EntityAccess.entity_id == entity_id)
    ).all():
        if a.user_id in investor_ids:
            session.delete(a)
            access += 1

    return {"statements": statements, "access_grants": access}


def _slug_username(name: str) -> str:
    base = re.sub(r"[^a-z0-9]+", "", name.lower())
    return base or "investor"


def _as_float(v) -> float | None:
    if v is None:
        return None
    if isinstance(v, (int, float)):
        return float(v)
    try:
        return float(str(v).replace(",", "").replace("$", "").strip())
    except ValueError:
        return None


# --- ALLOC SI (eNAV investor roster) ---

def _parse_alloc_si(wb):
    """Return (as_of, investors) from an eNAV ALLOC SI sheet.

    investors: [{name, external_id, commitment, contributions, distributions, ending}].
    """
    ws = wb["ALLOC SI"]
    rows = list(ws.iter_rows(min_row=1, max_row=MAX_ROWS, max_col=MAX_COLS, values_only=True))

    header_idx = None
    header: list[str] = []
    for i, row in enumerate(rows):
        cells = [str(c).strip().upper() if isinstance(c, str) else "" for c in row]
        if "INVESTOR ID" in cells and "INVESTOR NAME" in cells:
            header_idx, header = i, cells
            break
    if header_idx is None:
        raise HTTPException(status_code=422, detail="Could not find the ALLOC SI header row.")

    def col(*names: str) -> int | None:
        for n in names:
            if n in header:
                return header.index(n)
        return None

    c_id = col("INVESTOR ID")
    c_type = col("INVESTOR TYPE")
    c_name = col("INVESTOR NAME")
    c_commit = col("COMMITTED CAPITAL")
    c_contrib = col("CONTRIBUTIONS")
    c_distrib = col("(DISTRIBUTIONS)", "DISTRIBUTIONS")
    c_ending = col("ENDING BALANCE")
    if c_ending is None:
        raise HTTPException(status_code=422, detail="ALLOC SI is missing an ENDING BALANCE column.")

    def amount(row, ci):
        if ci is None or ci >= len(row):
            return 0.0
        return _as_float(row[ci]) or 0.0

    investors: list[dict] = []
    for r in range(header_idx + 1, len(rows)):
        row = rows[r]
        name_raw = row[c_name] if c_name is not None and c_name < len(row) else None
        if not (isinstance(name_raw, str) and name_raw.strip()):
            continue
        nm = name_raw.strip()
        if nm.upper() in ("GP", "LP", "TOTAL"):
            continue
        itype = row[c_type] if c_type is not None and c_type < len(row) else None
        if itype is not None and str(itype).strip().upper() not in ("LP", ""):
            continue  # investors only (skip the GP entity)

        display = None
        if c_name is not None and c_name + 1 < len(row) and isinstance(row[c_name + 1], str) and row[c_name + 1].strip():
            display = row[c_name + 1].strip()
        name = display or nm.title()

        inv_id = row[c_id] if c_id is not None and c_id < len(row) else None
        external_id = str(inv_id).strip() if inv_id not in (None, "") else None

        investors.append({
            "name": name,
            "external_id": external_id,
            "commitment": amount(row, c_commit),
            "contributions": amount(row, c_contrib),
            "distributions": abs(amount(row, c_distrib)),  # sheet may show as a credit
            "ending": amount(row, c_ending),
        })

    return _enav_as_of(wb), investors


# --- Subaccounts (names-across-columns) ---

def _parse_grid(wb):
    ws = wb[wb.sheetnames[0]]
    grid: list[list] = []
    for row in ws.iter_rows(min_row=1, max_row=MAX_ROWS, max_col=MAX_COLS, values_only=True):
        grid.append(list(row))

    header_idx = None
    total_col = None
    for i, row in enumerate(grid):
        for j, cell in enumerate(row):
            if isinstance(cell, str) and cell.strip().lower() == "total":
                header_idx, total_col = i, j
                break
        if header_idx is not None:
            break
    if header_idx is None:
        raise HTTPException(
            status_code=422,
            detail='Could not find a header row with a "Total" column. Check the spreadsheet layout.',
        )

    header = grid[header_idx]
    investor_cols = [
        j for j in range(len(header))
        if j < total_col
        and isinstance(header[j], str)
        and header[j].strip().lower() not in ("", "total")
    ]
    if not investor_cols:
        raise HTTPException(status_code=422, detail="No investor name columns found left of the Total column.")

    as_of: date | None = None
    for row in grid[: header_idx + 2]:
        for cell in row:
            if isinstance(cell, (datetime, date)):
                as_of = cell.date() if isinstance(cell, datetime) else cell
                break
        if as_of:
            break

    value_rows: list[ImportValueRow] = []
    for i in range(header_idx + 1, len(grid)):
        row = grid[i]
        label = row[0]
        if not (isinstance(label, str) and label.strip()):
            continue
        if any(_as_float(row[j]) is not None for j in investor_cols if j < len(row)):
            value_rows.append(ImportValueRow(row_index=i, label=label.strip()))
    if not value_rows:
        raise HTTPException(status_code=422, detail="No value rows found under the header.")

    return header_idx, total_col, investor_cols, as_of, value_rows, grid


@router.post("/preview")
def preview_import(
    file: UploadFile = File(...),
    entity_id: int | None = Form(None),
    row_index: int | None = Form(None),
    password: str | None = Form(None),
    admin: User = Depends(require_internal_admin),
    session: Session = Depends(get_session),
) -> CapitalImportPreview:
    if entity_id is not None and session.get(Entity, entity_id) is None:
        raise HTTPException(status_code=404, detail="Entity not found")

    try:
        file_bytes = storage.read_capped(file)
    except storage.UploadTooLarge as exc:
        raise HTTPException(status_code=413, detail=str(exc))
    wb = _open_workbook(file_bytes, password)

    investors = session.exec(select(User).where(User.role == UserRole.investor)).all()
    by_name = {u.name.strip().lower(): u for u in investors}
    by_username = {u.username.strip().lower(): u for u in investors}
    by_extid = {u.external_investor_id: u for u in investors if u.external_investor_id}

    def match_for(name: str, external_id: str | None):
        if external_id and external_id in by_extid:
            return by_extid[external_id]
        return by_name.get(name.strip().lower()) or by_username.get(name.strip().lower())

    if "ALLOC SI" in wb.sheetnames:
        as_of, roster = _parse_alloc_si(wb)
        previews: list[ImportInvestorPreview] = []
        for inv in roster:
            m = match_for(inv["name"], inv["external_id"])
            previews.append(ImportInvestorPreview(
                source_name=inv["name"],
                column_index=0,
                value_dollars=inv["ending"],
                commitment_dollars=inv["commitment"],
                contributions_dollars=inv["contributions"],
                distributions_dollars=inv["distributions"],
                external_id=inv["external_id"],
                matched_user_id=m.id if m else None,
                matched_username=m.username if m else None,
                suggested_username=None if m else _slug_username(inv["name"]),
            ))
        # Roster mode: no per-row value picker needed.
        return CapitalImportPreview(as_of_date=as_of, value_rows=[], chosen_row_index=0, investors=previews)

    # Fallback: subaccounts layout (single value figure per investor).
    header_idx, total_col, investor_cols, as_of, value_rows, grid = _parse_grid(wb)
    valid_indexes = {vr.row_index for vr in value_rows}
    chosen = row_index if row_index in valid_indexes else value_rows[0].row_index
    chosen_row = grid[chosen]

    previews = []
    for col_i in investor_cols:
        source_name = str(grid[header_idx][col_i]).strip()
        value = _as_float(chosen_row[col_i]) if col_i < len(chosen_row) else None
        m = match_for(source_name, None)
        previews.append(ImportInvestorPreview(
            source_name=source_name,
            column_index=col_i,
            value_dollars=value or 0.0,
            matched_user_id=m.id if m else None,
            matched_username=m.username if m else None,
            suggested_username=None if m else _slug_username(source_name),
        ))
    return CapitalImportPreview(as_of_date=as_of, value_rows=value_rows, chosen_row_index=chosen, investors=previews)


@router.post("/commit")
def commit_import(
    body: CapitalImportCommit,
    admin: User = Depends(require_internal_admin),
    session: Session = Depends(get_session),
) -> dict:
    entity = session.get(Entity, body.entity_id)
    if entity is None:
        raise HTTPException(status_code=404, detail="Entity not found")

    created_accounts = 0
    updated_accounts = 0
    statements = 0

    for inv in body.investors:
        if inv.action == "skip":
            continue

        if inv.action == "create":
            if not inv.username or not inv.name:
                raise HTTPException(
                    status_code=400,
                    detail=f"New member '{inv.name or inv.username}' needs a name and username.",
                )
            if session.exec(select(User).where(User.username == inv.username)).first():
                raise HTTPException(status_code=409, detail=f"Username '{inv.username}' already taken.")
            if inv.email and session.exec(select(User).where(User.email == inv.email)).first():
                raise HTTPException(status_code=409, detail=f"Email '{inv.email}' already in use.")
            # New members start on the shared default password (login enabled) so the admin
            # can send credentials right away; each investor rotates it in the portal.
            pw = inv.password or config.DEFAULT_INVESTOR_PASSWORD
            user = User(
                name=inv.name,
                username=inv.username,
                email=inv.email or None,
                password_hash=hash_password(pw),
                role=UserRole.investor,
                login_enabled=True,
                external_investor_id=inv.external_id,
                # First login forces a change while they're on the shared default.
                must_change_password=(pw == config.DEFAULT_INVESTOR_PASSWORD),
            )
            session.add(user)
            session.flush()
            created_accounts += 1
        elif inv.action == "match":
            user = session.get(User, inv.user_id) if inv.user_id else None
            if user is None or user.role != UserRole.investor:
                raise HTTPException(status_code=400, detail="match requires a valid investor user_id.")
            if inv.external_id and not user.external_investor_id:
                user.external_investor_id = inv.external_id
                session.add(user)
        else:
            raise HTTPException(status_code=400, detail=f"Unknown action '{inv.action}'.")

        has_access = session.exec(
            select(EntityAccess).where(
                EntityAccess.user_id == user.id, EntityAccess.entity_id == body.entity_id
            )
        ).first()
        if has_access is None:
            session.add(EntityAccess(user_id=user.id, entity_id=body.entity_id))

        cents = lambda d: round(d * 100)
        existing = session.exec(
            select(CapitalAccountStatement).where(
                CapitalAccountStatement.entity_id == body.entity_id,
                CapitalAccountStatement.investor_user_id == user.id,
                CapitalAccountStatement.as_of_date == body.as_of_date,
            )
        ).first()
        if existing:
            existing.commitment_cents = cents(inv.commitment_dollars)
            existing.contributions_cents = cents(inv.contributions_dollars)
            existing.distributions_cents = cents(inv.distributions_dollars)
            existing.ending_balance_cents = cents(inv.value_dollars)
            session.add(existing)
            updated_accounts += 1
        else:
            session.add(CapitalAccountStatement(
                entity_id=body.entity_id,
                investor_user_id=user.id,
                as_of_date=body.as_of_date,
                commitment_cents=cents(inv.commitment_dollars),
                contributions_cents=cents(inv.contributions_dollars),
                distributions_cents=cents(inv.distributions_dollars),
                ending_balance_cents=cents(inv.value_dollars),
            ))
        statements += 1

    record_audit(session, admin.id, "import", "capital_account", body.entity_id, {
        "as_of_date": str(body.as_of_date),
        "created_accounts": created_accounts,
        "statements": statements,
    })
    session.commit()
    return {
        "status": "ok",
        "created_accounts": created_accounts,
        "matched_accounts_updated": updated_accounts,
        "statements_written": statements,
    }


@router.post("/batch")
def batch_import(
    files: list[UploadFile] = File(...),
    entity_id: int = Form(...),
    password: str | None = Form(None),
    admin: User = Depends(require_internal_admin),
    session: Session = Depends(get_session),
) -> BatchCapitalImportResult:
    """Backfill several quarters of capital history for one fund from a batch of eNAV files.

    Each file is a full eNAV workbook; its ALLOC SI roster is auto-matched (by fund-admin
    investor ID, else name/username) against existing members and each member's capital
    statement is upserted at the file's own as-of date — so older files add historical points
    without touching the latest. Members not already in the system are skipped and reported
    (no account creation). One file failing (bad password, no ALLOC SI, unreadable date) is
    reported per-file and does not abort the rest.
    """
    entity = session.get(Entity, entity_id)
    if entity is None:
        raise HTTPException(status_code=404, detail="Entity not found")

    investors = session.exec(select(User).where(User.role == UserRole.investor)).all()
    by_name = {u.name.strip().lower(): u for u in investors}
    by_username = {u.username.strip().lower(): u for u in investors}
    by_extid = {u.external_investor_id: u for u in investors if u.external_investor_id}

    def match_for(name: str, external_id: str | None) -> User | None:
        if external_id and external_id in by_extid:
            return by_extid[external_id]
        return by_name.get(name.strip().lower()) or by_username.get(name.strip().lower())

    cents = lambda d: round((d or 0.0) * 100)
    results: list[BatchCapitalFileResult] = []
    total_statements = 0

    for upload in files:
        fname = upload.filename or "(unnamed)"
        res = BatchCapitalFileResult(filename=fname)
        try:
            wb = _open_workbook(storage.read_capped(upload), password)
            if "ALLOC SI" not in wb.sheetnames:
                raise HTTPException(status_code=422, detail="No ALLOC SI tab found in this workbook.")
            as_of, roster = _parse_alloc_si(wb)
            if as_of is None:
                raise HTTPException(status_code=422, detail="Could not determine the as-of date from the file.")
            res.as_of_date = as_of

            for inv in roster:
                user = match_for(inv["name"], inv["external_id"])
                if user is None:
                    res.skipped.append(inv["name"])
                    continue

                # Record the fund-admin ID on first sighting so later files match by ID too.
                if inv["external_id"] and not user.external_investor_id:
                    user.external_investor_id = inv["external_id"]
                    by_extid[inv["external_id"]] = user
                    session.add(user)

                has_access = session.exec(
                    select(EntityAccess).where(
                        EntityAccess.user_id == user.id, EntityAccess.entity_id == entity_id
                    )
                ).first()
                if has_access is None:
                    session.add(EntityAccess(user_id=user.id, entity_id=entity_id))

                existing = session.exec(
                    select(CapitalAccountStatement).where(
                        CapitalAccountStatement.entity_id == entity_id,
                        CapitalAccountStatement.investor_user_id == user.id,
                        CapitalAccountStatement.as_of_date == as_of,
                    )
                ).first()
                if existing:
                    existing.commitment_cents = cents(inv["commitment"])
                    existing.contributions_cents = cents(inv["contributions"])
                    existing.distributions_cents = cents(inv["distributions"])
                    existing.ending_balance_cents = cents(inv["ending"])
                    session.add(existing)
                    res.updated += 1
                else:
                    session.add(CapitalAccountStatement(
                        entity_id=entity_id,
                        investor_user_id=user.id,
                        as_of_date=as_of,
                        commitment_cents=cents(inv["commitment"]),
                        contributions_cents=cents(inv["contributions"]),
                        distributions_cents=cents(inv["distributions"]),
                        ending_balance_cents=cents(inv["ending"]),
                    ))
                res.matched += 1
                res.statements_written += 1

            record_audit(session, admin.id, "import_batch", "capital_account", entity_id, {
                "file": fname,
                "as_of_date": str(as_of),
                "matched": res.matched,
                "skipped": len(res.skipped),
            })
            session.commit()
            total_statements += res.statements_written
        except HTTPException as e:
            # Our own controlled messages (bad password, no ALLOC SI, no date) are safe to show.
            session.rollback()
            res.error = e.detail
        except storage.UploadTooLarge:
            session.rollback()
            res.error = "File is too large."
        except Exception:  # noqa: BLE001 — keep going on any parse failure; don't leak internals
            session.rollback()
            res.error = "Could not process this file (unexpected format or error)."
        results.append(res)

    return BatchCapitalImportResult(
        entity_id=entity_id, files=results, total_statements=total_statements
    )