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

"""XLSX/CSV import of a fund's holdings and NAV from a fund-administrator eNAV pack.

Reads the "HLD" (Holdings Report) sheet of an administrator eNAV workbook: each
security row becomes a holding + position, and its market value (book) becomes the
valuation for the quarter. Encrypted workbooks are decrypted with the open password.
A plain holdings CSV is also accepted.
"""

import csv
import io
from datetime import date, datetime
from typing import Any

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

from ten31portal.audit import record_audit
from ten31portal.auth import require_role
from ten31portal.database import get_session
from ten31portal.models import (
    Entity, EntityType, Holding, Position,
    User, UserRole, Valuation, ValuationRound, RoundStatus,
)

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




def _parse_money(raw: str) -> int | None:
    """Parse dollar strings like '$3.3M', '$61.9M', '$1,234,567', '3300000' to cents."""
    if not raw or not raw.strip():
        return None
    s = raw.strip().replace(",", "").replace("$", "")
    multiplier = 1
    if s.upper().endswith("M"):
        multiplier = 1_000_000
        s = s[:-1]
    elif s.upper().endswith("K"):
        multiplier = 1_000
        s = s[:-1]
    elif s.upper().endswith("B"):
        multiplier = 1_000_000_000
        s = s[:-1]
    try:
        return round(float(s) * multiplier * 100)
    except ValueError:
        return None


def _parse_date(raw: str) -> date | None:
    """Try common date formats."""
    for fmt in ("%Y-%m-%d", "%m/%d/%Y", "%m/%d/%y", "%Y/%m/%d"):
        try:
            return datetime.strptime(raw.strip(), fmt).date()
        except ValueError:
            continue
    return None


def _dollars_to_cents(val: float | int) -> int:
    """Convert dollar amount to integer cents."""
    return round(float(val) * 100)


# --- eNAV holdings import (XLSX) ---

def _open_workbook(file_bytes: bytes, password: str | None):
    """Load an xlsx workbook, decrypting an encrypted (password-protected) file if needed."""
    # A normal .xlsx is a zip ("PK"); an encrypted Office file is an OLE2 container.
    if file_bytes[:2] == b"PK":
        return openpyxl.load_workbook(io.BytesIO(file_bytes), data_only=True)
    try:
        office = msoffcrypto.OfficeFile(io.BytesIO(file_bytes))
        office.load_key(password=password or "VelvetSweatshop")
        out = io.BytesIO()
        office.decrypt(out)
        out.seek(0)
        return openpyxl.load_workbook(out, data_only=True)
    except Exception:
        if not password:
            raise HTTPException(
                status_code=400,
                detail="This spreadsheet is password-protected. Enter the open password and try again.",
            )
        raise HTTPException(
            status_code=400,
            detail="Could not open the spreadsheet. The password may be incorrect.",
        )


def _enav_as_of(wb) -> date | None:
    """Find the report date, e.g. a 'DECEMBER 31, 2025' / 'AS OF ...' cell near the top."""
    for sheet in (["MENU", "HLD"] if "MENU" in wb.sheetnames else wb.sheetnames):
        ws = wb[sheet]
        for row in ws.iter_rows(min_row=1, max_row=6, max_col=2, values_only=True):
            for cell in row:
                if isinstance(cell, datetime):
                    return cell.date()
                if isinstance(cell, date):
                    return cell
                if isinstance(cell, str):
                    text = cell.replace("AS OF", "").strip()
                    for fmt in ("%B %d, %Y", "%b %d, %Y", "%m/%d/%Y"):
                        try:
                            return datetime.strptime(text, fmt).date()
                        except ValueError:
                            continue
    return None


def _parse_schedule_xlsx(file_bytes: bytes, password: str | None = None) -> tuple[str | None, date | None, list[dict], list[dict], list[dict]]:
    """
    Parse the HLD (Holdings Report) sheet of a fund-administrator eNAV workbook.

    Each security row becomes a holding (issuer) + position (security), with cost basis
    and market value (book). Returns: (entity_name, as_of, holdings, positions, errors).
    """
    wb = _open_workbook(file_bytes, password)

    entity_name: str | None = None
    holdings_preview: list[dict] = []
    positions_preview: list[dict] = []
    errors: list[dict] = []

    # Fund name from the MENU cover sheet (row 2) when present.
    if "MENU" in wb.sheetnames:
        v = wb["MENU"].cell(row=2, column=1).value
        if v:
            entity_name = str(v).strip()

    as_of = _enav_as_of(wb)

    if "HLD" not in wb.sheetnames:
        errors.append({"row": 0, "errors": [
            "No 'HLD' (Holdings Report) sheet found. This does not look like an eNAV workbook."
        ]})
        return entity_name, as_of, holdings_preview, positions_preview, errors

    ws = wb["HLD"]
    rows = list(ws.iter_rows(min_row=1, max_row=ws.max_row, max_col=20, values_only=True))

    # Locate the header row (the one whose first cell is "SECURITY NAME").
    header_idx = None
    for i, row in enumerate(rows):
        if row and isinstance(row[0], str) and row[0].strip().upper() == "SECURITY NAME":
            header_idx = i
            break
    if header_idx is None:
        errors.append({"row": 0, "errors": ["Could not find the holdings header row in the HLD sheet."]})
        return entity_name, as_of, holdings_preview, positions_preview, errors

    header = [str(c).strip().upper() if isinstance(c, str) else "" for c in rows[header_idx]]

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

    c_name = col("SECURITY NAME")
    c_qty = col("QUANTITY")
    c_cost = col("COST BASIS - BOOK", "COST BASIS - LOCAL")
    c_value = col("MARKET VALUE (BOOK)", "MARKET VALUE (LOCAL)", "MARKET VALUE - BOOK")

    if c_value is None or c_name is None:
        errors.append({"row": header_idx + 1, "errors": [
            "HLD sheet is missing a SECURITY NAME or MARKET VALUE column."
        ]})
        return entity_name, as_of, holdings_preview, positions_preview, errors

    seen_companies: set[str] = set()
    for r in range(header_idx + 1, len(rows)):
        row = rows[r]
        raw_name = row[c_name] if c_name < len(row) else None
        if raw_name is None or not str(raw_name).strip():
            continue
        name = str(raw_name).strip()
        # Stop at total / report-total summary rows.
        if name.upper().startswith("TOTAL") or name.upper().startswith("REPORT TOTAL"):
            continue

        # Group tranches of the same issuer: holding = issuer (before " - "); position = full name.
        company = name.split(" - ")[0].strip() or name
        security = name

        if company not in seen_companies:
            holdings_preview.append({"row": r + 1, "company_name": company})
            seen_companies.add(company)

        # Quantity -> shares (skip non-numeric like "N/A").
        shares: str | None = None
        if c_qty is not None and c_qty < len(row) and row[c_qty] is not None:
            qv = row[c_qty]
            if isinstance(qv, (int, float)) and qv != 0:
                shares = f"{qv:g}"

        cost_cents = None
        if c_cost is not None and c_cost < len(row) and isinstance(row[c_cost], (int, float)):
            cost_cents = _dollars_to_cents(row[c_cost])

        value_cents = None
        if isinstance(row[c_value], (int, float)):
            value_cents = _dollars_to_cents(row[c_value])

        positions_preview.append({
            "row": r + 1,
            "company_name": company,
            "security_name": security,
            "investment_date": None,  # eNAV holdings report has no acquisition date
            "shares": shares,
            "cost_cents": cost_cents,
            "value_cents": value_cents,
            "valuation_date": str(as_of) if as_of else None,
        })

    return entity_name, as_of, holdings_preview, positions_preview, errors




def reset_entity_holdings(entity_id: int, session: Session) -> dict[str, int]:
    """Delete all holdings, positions, valuations, and rounds for one entity.

    For a clean restart — e.g. after switching the source workbook (Carta → eNAV) renamed every
    position, so the old and new rows can't be matched and both get counted. Capital-account
    statements (investor data) are left untouched. The caller should re-import afterwards.
    """
    rounds = session.exec(
        select(ValuationRound).where(ValuationRound.entity_id == entity_id)
    ).all()
    round_ids = [r.id for r in rounds]
    valuations = 0
    if round_ids:
        for v in session.exec(
            select(Valuation).where(col(Valuation.round_id).in_(round_ids))
        ).all():
            session.delete(v)
            valuations += 1

    holdings = session.exec(select(Holding).where(Holding.entity_id == entity_id)).all()
    holding_ids = [h.id for h in holdings]
    positions = 0
    if holding_ids:
        for p in session.exec(
            select(Position).where(col(Position.holding_id).in_(holding_ids))
        ).all():
            session.delete(p)
            positions += 1

    for r in rounds:
        session.delete(r)
    for h in holdings:
        session.delete(h)
    session.flush()
    return {
        "rounds": len(rounds),
        "holdings": len(holdings),
        "positions": positions,
        "valuations": valuations,
    }


def dedupe_entity(entity_id: int, session: Session) -> dict[str, int]:
    """Collapse exact-duplicate holdings and positions for one entity.

    Repeated imports on older builds created a second copy of each holding/position, which
    inflated the entity's "Invested" total (the rollup sums cost across every position). This
    keeps the lowest-id copy, re-points its positions/valuations, removes duplicate valuations
    in the same round, and deletes the leftovers. Safe to run repeatedly (a no-op once clean).
    """
    removed_holdings = 0
    removed_positions = 0

    # 1) Merge holdings with the same name (case-insensitive) into the lowest-id one.
    holdings = session.exec(
        select(Holding).where(Holding.entity_id == entity_id).order_by(Holding.id)  # type: ignore[arg-type]
    ).all()
    keep_by_name: dict[str, Holding] = {}
    for h in holdings:
        key = h.company_name.strip().lower()
        keeper = keep_by_name.get(key)
        if keeper is None:
            keep_by_name[key] = h
            continue
        for p in session.exec(select(Position).where(Position.holding_id == h.id)).all():
            p.holding_id = keeper.id
            session.add(p)
        session.flush()
        session.delete(h)
        removed_holdings += 1
    session.flush()

    # 2) Within each surviving holding, merge positions with the same security name.
    for keeper_holding in keep_by_name.values():
        positions = session.exec(
            select(Position)
            .where(Position.holding_id == keeper_holding.id)
            .order_by(Position.id)  # type: ignore[arg-type]
        ).all()
        keep_by_sec: dict[str, Position] = {}
        for p in positions:
            key = p.security_name.strip().lower()
            keeper = keep_by_sec.get(key)
            if keeper is None:
                keep_by_sec[key] = p
                continue
            # Move this duplicate's valuations onto the keeper, dropping any that would collide
            # with an existing valuation for the same round (that collision IS the double-count).
            for v in session.exec(select(Valuation).where(Valuation.position_id == p.id)).all():
                clash = session.exec(
                    select(Valuation).where(
                        Valuation.round_id == v.round_id,
                        Valuation.position_id == keeper.id,
                    )
                ).first()
                if clash is not None:
                    session.delete(v)
                else:
                    v.position_id = keeper.id
                    session.add(v)
            session.flush()
            session.delete(p)
            removed_positions += 1
    session.flush()

    return {"removed_holdings": removed_holdings, "removed_positions": removed_positions}


@router.post("/schedule")
def import_schedule(
    file: UploadFile = File(...),
    as_of: date | None = Query(default=None),
    entity_id: int | None = Query(default=None),
    create_entity_type: EntityType = Query(default=EntityType.fund),
    create_vintage_year: int | None = Query(default=None),
    commit: bool = Query(default=True),
    replace_existing: bool = Query(default=False),
    password: str | None = Form(default=None),
    user: User = Depends(require_role(UserRole.approver, UserRole.cfo, UserRole.operations)),
    session: Session = Depends(get_session),
) -> dict[str, Any]:
    file_bytes = file.file.read()
    filename = file.filename or ""

    # Parse the file into holdings/positions previews.
    if filename.lower().endswith((".xlsx", ".xls")):
        source_entity_name, parsed_as_of, holdings_preview, positions_preview, errors = (
            _parse_schedule_xlsx(file_bytes, password)
        )
    else:
        source_entity_name = None
        parsed_as_of = None
        holdings_preview, positions_preview, errors = _parse_schedule_csv(file_bytes)

    # As-of date: use the explicit value, else the date read from the sheet.
    as_of = as_of or parsed_as_of
    if as_of is None:
        raise HTTPException(
            status_code=400,
            detail="Could not determine the quarter-end date. Set the as-of date and try again.",
        )

    # Resolve entity
    entity: Entity | None = None
    entity_resolution: str = "existing"  # "existing" | "matched" | "will_create"

    if entity_id is not None:
        # Explicit entity_id: use it directly
        entity = session.get(Entity, entity_id)
        if entity is None:
            raise HTTPException(status_code=404, detail="Entity not found")
        entity_resolution = "existing"
    elif source_entity_name:
        # Try matching by name from row 1
        entity = session.exec(
            select(Entity).where(Entity.name == source_entity_name)
        ).first()
        if entity:
            entity_resolution = "matched"
        else:
            entity_resolution = "will_create"
    else:
        raise HTTPException(
            status_code=400,
            detail="No entity selected and no fund name found in the file. Choose the fund to import into.",
        )

    # Dry-run: report resolution without writing
    if not commit:
        resolution_info: dict[str, Any] = {"resolution": entity_resolution}
        if entity:
            resolution_info["entity_id"] = entity.id
            resolution_info["entity_name"] = entity.name
            resolution_info["entity_type"] = entity.type.value
        else:
            resolution_info["entity_name"] = source_entity_name
            resolution_info["entity_type"] = create_entity_type.value
            resolution_info["vintage_year"] = create_vintage_year
        return {
            "committed": False,
            "source_entity_name": source_entity_name,
            "entity": resolution_info,
            "holdings": holdings_preview,
            "positions": positions_preview,
            "errors": errors,
            "seed_round": {"entity_id": entity.id if entity else None, "quarter_end": str(as_of)},
        }

    # Commit path: create entity if needed
    if entity_resolution == "will_create" and source_entity_name:
        entity = Entity(
            name=source_entity_name,
            type=create_entity_type,
            vintage_year=create_vintage_year,
        )
        session.add(entity)
        session.flush()
        record_audit(session, user.id, "create", "entity", entity.id, {
            "name": source_entity_name,
            "type": create_entity_type.value,
            "source": "schedule_import",
        })

    assert entity is not None
    resolved_entity_id = entity.id

    # Replace mode: wipe the fund's existing holdings/positions/rounds first so the file becomes
    # the single source of truth. Use this after a source change (e.g. Carta → eNAV) renamed the
    # positions, leaving un-matchable rows that inflate the totals.
    if replace_existing:
        reset_entity_holdings(resolved_entity_id, session)

    # Self-heal any duplicate holdings/positions left by imports on older builds, so this
    # import's upsert lands on a single clean copy and "Invested" stops double-counting.
    dedupe_entity(resolved_entity_id, session)

    # An existing round at this quarter: a NAV re-import should UPDATE it in place (refresh
    # to the latest file) instead of stacking a second round and doubling the totals. Only
    # import-created seed rounds are refreshable; a manually-signed valuation round is left
    # protected.
    existing_round = session.exec(
        select(ValuationRound).where(
            ValuationRound.entity_id == resolved_entity_id,
            ValuationRound.quarter_end == as_of,
        )
    ).first()
    reused_round = False
    seed_round: ValuationRound | None = None
    if existing_round is not None:
        if not existing_round.is_seed:
            raise HTTPException(
                status_code=409,
                detail=f"A signed valuation round already exists for this entity and quarter ({as_of}). Cannot overwrite signed history.",
            )
        seed_round = existing_round
        reused_round = True

    # Commit: create holdings, positions, seed round
    holding_map: dict[str, Holding] = {}

    for hp in holdings_preview:
        name = hp["company_name"]
        existing = session.exec(
            select(Holding).where(
                Holding.entity_id == resolved_entity_id,
                Holding.company_name == name,
            )
        ).first()
        if existing:
            holding_map[name] = existing
        else:
            h = Holding(entity_id=resolved_entity_id, company_name=name)
            session.add(h)
            session.flush()
            holding_map[name] = h

    # Create the seed round, or reuse the existing one for this quarter (re-import).
    if seed_round is None:
        seed_round = ValuationRound(
            entity_id=resolved_entity_id,
            quarter_end=as_of,
            status=RoundStatus.approved,
            is_seed=True,
            approved_by=user.id,
            approved_at=datetime.utcnow(),
        )
        session.add(seed_round)
        session.flush()

    positions_created = 0
    positions_updated = 0
    seen_position_ids: set[int] = set()
    for pp in positions_preview:
        company = pp["company_name"]
        holding = holding_map.get(company)
        if holding is None:
            existing = session.exec(
                select(Holding).where(
                    Holding.entity_id == resolved_entity_id,
                    Holding.company_name == company,
                )
            ).first()
            if existing:
                holding = existing
            else:
                holding = Holding(entity_id=resolved_entity_id, company_name=company)
                session.add(holding)
                session.flush()
            holding_map[company] = holding

        # Upsert position by (holding, security_name)
        pos = session.exec(
            select(Position).where(
                Position.holding_id == holding.id,
                Position.security_name == pp["security_name"],
            )
        ).first()
        if pos is None:
            pos = Position(
                holding_id=holding.id,
                security_name=pp["security_name"],
                investment_date=_parse_date(pp["investment_date"]) if pp["investment_date"] else as_of,
                shares=pp["shares"],
                cost_cents=pp["cost_cents"] or 0,
            )
            session.add(pos)
            session.flush()
            positions_created += 1
        else:
            if pp["investment_date"]:
                pos.investment_date = _parse_date(pp["investment_date"])
            if pp["shares"]:
                pos.shares = pp["shares"]
            if pp["cost_cents"] is not None:
                pos.cost_cents = pp["cost_cents"]
            session.add(pos)
            session.flush()
            positions_updated += 1

        # Upsert this position's valuation in the round, so a re-import refreshes the value
        # in place instead of adding a second one (which would double the quarter's NAV).
        val = session.exec(
            select(Valuation).where(
                Valuation.round_id == seed_round.id,
                Valuation.position_id == pos.id,
            )
        ).first()
        if val is None:
            session.add(Valuation(
                round_id=seed_round.id,
                position_id=pos.id,
                value_cents=pp["value_cents"] or 0,
            ))
        else:
            val.value_cents = pp["value_cents"] or 0
            session.add(val)
        seen_position_ids.add(pos.id)

    # On re-import, drop valuations for holdings no longer in the file so the quarter's NAV
    # equals the new file's total (no leftovers from the prior import).
    if reused_round:
        stale_q = select(Valuation).where(Valuation.round_id == seed_round.id)
        if seen_position_ids:
            stale_q = stale_q.where(col(Valuation.position_id).not_in(seen_position_ids))
        for stale_val in session.exec(stale_q).all():
            session.delete(stale_val)

    record_audit(session, user.id, "import_schedule", "entity", resolved_entity_id, {
        "source_entity_name": source_entity_name,
        "entity_resolution": entity_resolution,
        "replaced_existing": replace_existing,
        "round_updated": reused_round,
        "holdings": len(holding_map),
        "positions_created": positions_created,
        "positions_updated": positions_updated,
        "seed_quarter": str(as_of),
    })
    session.commit()

    return {
        "committed": True,
        "source_entity_name": source_entity_name,
        "entity": {
            "resolution": entity_resolution,
            "entity_id": resolved_entity_id,
            "entity_name": entity.name,
            "entity_type": entity.type.value,
        },
        "holdings_count": len(holding_map),
        "positions_count": len(positions_preview),
        "positions_created": positions_created,
        "positions_updated": positions_updated,
        "seed_round_id": seed_round.id,
        "round_updated": reused_round,
        "replaced_existing": replace_existing,
        "errors": errors,
    }


def _parse_schedule_csv(file_bytes: bytes) -> tuple[list[dict], list[dict], list[dict]]:
    """Parse a plain holdings CSV (generic columns or an eNAV HLD export)."""
    content = file_bytes.decode("utf-8-sig")
    reader = csv.DictReader(io.StringIO(content))

    CSV_COLUMN_MAP = {
        "Company": "company_name",
        "Investment": "company_name",
        "Issuer": "company_name",
        "Security": "security_name",
        "Asset": "security_name",
        "Security Name": "security_name",
        "Investment Date": "investment_date",
        "Investment date": "investment_date",
        "Shares": "shares",
        "Quantity": "shares",
        "Cost": "cost_cents",
        "Cost Basis - Book": "cost_cents",
        "Cost Basis - Local": "cost_cents",
        "Value": "value_cents",
        "Market Value (Book)": "value_cents",
        "Market Value (Local)": "value_cents",
        "Market Value - Book": "value_cents",
    }

    holdings_preview: list[dict] = []
    positions_preview: list[dict] = []
    errors: list[dict] = []
    current_company: str | None = None

    for i, row in enumerate(reader, start=2):
        mapped: dict[str, Any] = {}
        for csv_col, model_field in CSV_COLUMN_MAP.items():
            val = row.get(csv_col)
            if val is None:
                for k, v in row.items():
                    if k.strip().lower() == csv_col.lower():
                        val = v
                        break
            if val is not None and model_field not in mapped:
                mapped[model_field] = (val or "").strip()

        company = mapped.get("company_name", "").strip()
        security = mapped.get("security_name", "").strip()
        # Skip total / summary rows.
        if security.upper().startswith("TOTAL") or company.upper().startswith("TOTAL"):
            continue
        # eNAV-style rows have only a security name; derive the issuer from its prefix.
        if security and not company:
            company = security.split(" - ")[0].strip()
            mapped["company_name"] = company

        if company and not security:
            current_company = company
            holdings_preview.append({"row": i, "company_name": company})
            continue

        if not company and current_company:
            company = current_company
        elif company:
            if company not in [h["company_name"] for h in holdings_preview]:
                holdings_preview.append({"row": i, "company_name": company})
            current_company = company

        if not security:
            continue

        row_errors: list[str] = []
        inv_date = None
        if mapped.get("investment_date"):
            inv_date = _parse_date(mapped["investment_date"])
            if inv_date is None:
                row_errors.append(f"Cannot parse date: {mapped['investment_date']}")

        shares = None
        raw_shares = mapped.get("shares", "")
        if raw_shares:
            s = raw_shares.replace(",", "")
            try:
                sv = float(s)
                if sv != 0:
                    shares = s
            except ValueError:
                row_errors.append(f"Cannot parse shares: {raw_shares}")

        cost_cents = None
        if mapped.get("cost_cents"):
            cost_cents = _parse_money(mapped["cost_cents"])
            if cost_cents is None:
                row_errors.append(f"Cannot parse cost: {mapped['cost_cents']}")

        value_cents = None
        if mapped.get("value_cents"):
            value_cents = _parse_money(mapped["value_cents"])
            if value_cents is None:
                row_errors.append(f"Cannot parse value: {mapped['value_cents']}")

        if row_errors:
            errors.append({"row": i, "errors": row_errors, "raw": dict(row)})
            continue

        positions_preview.append({
            "row": i,
            "company_name": company,
            "security_name": security,
            "investment_date": str(inv_date) if inv_date else None,
            "shares": shares,
            "cost_cents": cost_cents,
            "value_cents": value_cents,
        })

    return holdings_preview, positions_preview, errors