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

"""Entity CRUD and rollup endpoints."""

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy import func, literal
from sqlmodel import Session, col, select

from ten31portal.audit import record_audit
from ten31portal.auth import (
    accessible_entity_ids, get_current_user, household_user_ids, require_internal,
    require_writer,
)
from ten31portal.database import get_session
from ten31portal.models import (
    CapitalAccountStatement, Entity, EntityAccess, EntityStake, EntityStatus, Holding,
    Position, UserRole, Valuation, ValuationRound, RoundStatus, User,
)
from ten31portal.schemas import (
    AssetBalancesResponse, CapitalAccountResponse, EntityCreate, EntityResponse,
    EntityStakeCreate, EntityStakeResponse, EntityUpdate, PartnerExitUpdate,
    PartnerResponse,
)

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


class EntityRollupItem(BaseModel):
    id: int
    name: str
    type: str
    vintage_year: int | None
    fund_size_cents: int | None
    status: str
    committed_cents: int          # total LP commitments (latest per investor)
    invested_cents: int
    last_signed_value_cents: int


@router.get("/rollup")
def entity_rollup(
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> list[EntityRollupItem]:
    """Per-entity invested and last-signed-value in a single pass."""
    allowed = accessible_entity_ids(user, session)
    entities = session.exec(select(Entity)).all()
    if allowed is not None:
        entities = [e for e in entities if e.id in allowed]
    result: list[EntityRollupItem] = []

    for ent in entities:
        # Sum cost_cents across all positions in this entity's holdings
        invested_row = session.exec(
            select(func.coalesce(func.sum(Position.cost_cents), 0))
            .join(Holding, Position.holding_id == Holding.id)
            .where(Holding.entity_id == ent.id)
        ).one()
        invested_cents = int(invested_row)

        # Find the latest approved round for this entity
        latest_round = session.exec(
            select(ValuationRound)
            .where(
                ValuationRound.entity_id == ent.id,
                ValuationRound.status == RoundStatus.approved,
            )
            .order_by(ValuationRound.quarter_end.desc())  # type: ignore[union-attr]
        ).first()

        last_signed_value_cents = 0
        if latest_round:
            val_sum = session.exec(
                select(func.coalesce(func.sum(Valuation.value_cents), 0))
                .where(Valuation.round_id == latest_round.id)
            ).one()
            last_signed_value_cents = int(val_sum)

        # Total committed capital = each investor's most recent commitment for this entity.
        # Exited members (stake sold/transferred) are skipped — their buyer's commitment now
        # appears on the roster, so counting both would double the fund's committed total.
        stmts = session.exec(
            select(CapitalAccountStatement)
            .where(CapitalAccountStatement.entity_id == ent.id)
            .order_by(CapitalAccountStatement.as_of_date.desc())  # type: ignore[union-attr]
        ).all()
        exited_ids = set(session.exec(
            select(EntityAccess.user_id).where(
                EntityAccess.entity_id == ent.id,
                col(EntityAccess.exited_on).is_not(None),
            )
        ).all())
        committed_cents = 0
        seen_investors: set[int] = set(exited_ids)
        for st in stmts:
            if st.investor_user_id in seen_investors:
                continue
            seen_investors.add(st.investor_user_id)
            committed_cents += st.commitment_cents

        result.append(EntityRollupItem(
            id=ent.id,
            name=ent.name,
            type=ent.type.value,
            vintage_year=ent.vintage_year,
            fund_size_cents=ent.fund_size_cents,
            status=ent.status.value,
            committed_cents=committed_cents,
            invested_cents=invested_cents,
            last_signed_value_cents=last_signed_value_cents,
        ))

    return result


@router.get("/{entity_id}/partners")
def list_partners(
    entity_id: int,
    user: User = Depends(require_internal),
    session: Session = Depends(get_session),
) -> list[PartnerResponse]:
    """Members (investors) granted access to this entity, with their latest capital value."""
    if session.get(Entity, entity_id) is None:
        raise HTTPException(status_code=404, detail="Entity not found")

    members = session.exec(
        select(User, EntityAccess)
        .join(EntityAccess, EntityAccess.user_id == User.id)
        .where(EntityAccess.entity_id == entity_id, User.role == UserRole.investor)
        .order_by(User.name)  # type: ignore[arg-type]
    ).all()

    result: list[PartnerResponse] = []
    for m, access in members:
        stmts = session.exec(
            select(CapitalAccountStatement)
            .where(
                CapitalAccountStatement.entity_id == entity_id,
                CapitalAccountStatement.investor_user_id == m.id,
            )
            .order_by(CapitalAccountStatement.as_of_date.desc())  # type: ignore[union-attr]
        ).all()
        latest = stmts[0] if stmts else None
        result.append(PartnerResponse(
            user_id=m.id,
            name=m.name,
            username=m.username,
            external_investor_id=m.external_investor_id,
            is_active=m.is_active,
            login_enabled=m.login_enabled,
            latest_commitment_cents=latest.commitment_cents if latest else None,
            latest_contributions_cents=latest.contributions_cents if latest else None,
            latest_distributions_cents=latest.distributions_cents if latest else None,
            latest_value_cents=latest.ending_balance_cents if latest else None,
            latest_as_of=latest.as_of_date if latest else None,
            statements_count=len(stmts),
            exited_on=access.exited_on,
        ))
    return result


@router.put("/{entity_id}/partners/{user_id}/exited")
def set_partner_exited(
    entity_id: int,
    user_id: int,
    body: PartnerExitUpdate,
    admin: User = Depends(require_writer),
    session: Session = Depends(get_session),
) -> PartnerResponse:
    """Mark a member as exited from this fund (stake sold/transferred), or clear it.

    Their statements and documents stay; the portal shows an Exited badge instead of a
    phantom -100% and drops the position from portfolio and fund committed totals.
    """
    if session.get(Entity, entity_id) is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    member = session.get(User, user_id)
    if member is None or member.role != UserRole.investor:
        raise HTTPException(status_code=400, detail="Only investor members can be marked exited")

    access = session.exec(
        select(EntityAccess).where(
            EntityAccess.entity_id == entity_id, EntityAccess.user_id == user_id
        )
    ).first()
    if access is None:
        # A manually-entered investor may have statements without an access grant yet
        # (e.g. from the Capital Accounts screen). Marking them exited creates the roster
        # row with the flag set; clearing an exit that doesn't exist stays a 404.
        if body.exited_on is None:
            raise HTTPException(status_code=404, detail="That member has no access to this fund")
        access = EntityAccess(user_id=user_id, entity_id=entity_id)
        session.add(access)
        session.flush()

    access.exited_on = body.exited_on
    session.add(access)
    record_audit(session, admin.id, "set_exited", "entity_access", access.id, {
        "entity_id": entity_id,
        "user_id": user_id,
        "exited_on": str(body.exited_on) if body.exited_on else None,
    })
    session.commit()

    # Return the member's refreshed partner row for easy UI updates.
    return next(p for p in list_partners(entity_id, admin, session) if p.user_id == user_id)


@router.delete("/{entity_id}/partners")
def clear_partners(
    entity_id: int,
    user: User = Depends(require_writer),
    session: Session = Depends(get_session),
) -> dict[str, int]:
    """Remove all partners from this fund — deletes its capital-account statements and the
    investors' access grants, but keeps the investor accounts (they belong to other funds).
    For undoing a wrong members import. Holdings/NAV are not affected."""
    # Local import avoids a module-load cycle (capital_import_router imports import_router).
    from ten31portal.routers.capital_import_router import reset_entity_partners

    if session.get(Entity, entity_id) is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    res = reset_entity_partners(entity_id, session)
    record_audit(session, user.id, "clear_partners", "entity", entity_id, res)
    session.commit()
    return res


@router.get("")
def list_entities(
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> list[EntityResponse]:
    allowed = accessible_entity_ids(user, session)
    rows = session.exec(select(Entity)).all()
    if allowed is not None:
        rows = [r for r in rows if r.id in allowed]
    return [EntityResponse.model_validate(r, from_attributes=True) for r in rows]


@router.get("/{entity_id}")
def get_entity(
    entity_id: int,
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> EntityResponse:
    allowed = accessible_entity_ids(user, session)
    if allowed is not None and entity_id not in allowed:
        raise HTTPException(status_code=404, detail="Entity not found")
    entity = session.get(Entity, entity_id)
    if entity is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    return EntityResponse.model_validate(entity, from_attributes=True)


def _validate_linked_user(linked_user_id: int | None, session: Session) -> None:
    """A linked account (for a GP entity that is also an LP) must be an investor account."""
    if linked_user_id is None:
        return
    linked = session.get(User, linked_user_id)
    if linked is None or linked.role != UserRole.investor:
        raise HTTPException(status_code=400, detail="Linked account must be an investor account.")


@router.post("", status_code=201)
def create_entity(
    body: EntityCreate,
    user: User = Depends(require_writer),
    session: Session = Depends(get_session),
) -> EntityResponse:
    _validate_linked_user(body.linked_user_id, session)
    entity = Entity(**body.model_dump())
    session.add(entity)
    session.flush()
    record_audit(session, user.id, "create", "entity", entity.id, body.model_dump())
    session.commit()
    session.refresh(entity)
    return EntityResponse.model_validate(entity, from_attributes=True)


@router.patch("/{entity_id}")
def update_entity(
    entity_id: int,
    body: EntityUpdate,
    user: User = Depends(require_writer),
    session: Session = Depends(get_session),
) -> EntityResponse:
    entity = session.get(Entity, entity_id)
    if entity is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    changes = body.model_dump(exclude_unset=True)
    if "linked_user_id" in changes:
        _validate_linked_user(changes["linked_user_id"], session)
    for key, val in changes.items():
        setattr(entity, key, val)
    session.add(entity)
    session.flush()
    # mode="json" so date fields (close_date) serialize into the audit JSON column.
    record_audit(session, user.id, "update", "entity", entity.id,
                 body.model_dump(exclude_unset=True, mode="json"))
    session.commit()
    session.refresh(entity)
    return EntityResponse.model_validate(entity, from_attributes=True)


@router.get("/{entity_id}/asset-balances")
def asset_balances(
    entity_id: int,
    user: User = Depends(require_internal),
    session: Session = Depends(get_session),
) -> AssetBalancesResponse:
    """A GP/mgmt entity's assets = the linked account's capital balances across the funds.

    Household-aware: if the linked account has other legal names linked to it (as the eNAV
    often splits one LLC across names), their balances are included too.
    """
    entity = session.get(Entity, entity_id)
    if entity is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    if entity.linked_user_id is None:
        return AssetBalancesResponse()
    linked = session.get(User, entity.linked_user_id)
    if linked is None:
        return AssetBalancesResponse(linked_user_id=entity.linked_user_id)

    household = household_user_ids(linked, session)
    rows = session.exec(
        select(CapitalAccountStatement)
        .where(col(CapitalAccountStatement.investor_user_id).in_(household))
        .order_by(col(CapitalAccountStatement.as_of_date).desc())
    ).all()
    names = dict(session.exec(
        select(User.id, User.name).where(col(User.id).in_({r.investor_user_id for r in rows}))
    ).all()) if rows else {}
    balances: list[CapitalAccountResponse] = []
    for r in rows:
        d = CapitalAccountResponse.model_validate(r, from_attributes=True)
        d.investor_name = names.get(r.investor_user_id)
        balances.append(d)
    return AssetBalancesResponse(
        linked_user_id=linked.id, linked_name=linked.name, balances=balances,
    )


# --- Entity stakes: a GP/mgmt entity's interest in the funds it manages ---

def _stake_response(stake: EntityStake, funds: dict[int, Entity]) -> EntityStakeResponse:
    data = EntityStakeResponse.model_validate(stake, from_attributes=True)
    fund = funds.get(stake.fund_entity_id)
    if fund is not None:
        data.fund_name = fund.name
        data.fund_type = fund.type
    return data


@router.get("/{entity_id}/stakes")
def list_stakes(
    entity_id: int,
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> list[EntityStakeResponse]:
    """The funds this entity holds a stake in (e.g. a GP's interest in its funds)."""
    allowed = accessible_entity_ids(user, session)
    if allowed is not None and entity_id not in allowed:
        raise HTTPException(status_code=404, detail="Entity not found")
    if session.get(Entity, entity_id) is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    rows = session.exec(
        select(EntityStake).where(EntityStake.holder_entity_id == entity_id)
    ).all()
    funds = {
        f.id: f for f in session.exec(
            select(Entity).where(col(Entity.id).in_({r.fund_entity_id for r in rows}))
        ).all()
    } if rows else {}
    return [_stake_response(r, funds) for r in rows]


@router.post("/{entity_id}/stakes", status_code=201)
def create_stake(
    entity_id: int,
    body: EntityStakeCreate,
    user: User = Depends(require_writer),
    session: Session = Depends(get_session),
) -> EntityStakeResponse:
    if session.get(Entity, entity_id) is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    if body.fund_entity_id == entity_id:
        raise HTTPException(status_code=400, detail="An entity cannot hold a stake in itself.")
    fund = session.get(Entity, body.fund_entity_id)
    if fund is None:
        raise HTTPException(status_code=404, detail="Fund not found")
    if session.exec(
        select(EntityStake).where(
            EntityStake.holder_entity_id == entity_id,
            EntityStake.fund_entity_id == body.fund_entity_id,
        )
    ).first():
        raise HTTPException(status_code=409, detail="A stake in this fund already exists.")

    stake = EntityStake(
        holder_entity_id=entity_id,
        fund_entity_id=body.fund_entity_id,
        ownership_pct=body.ownership_pct,
        value_cents=round(body.value_dollars * 100) if body.value_dollars is not None else None,
        note=body.note,
    )
    session.add(stake)
    session.flush()
    record_audit(session, user.id, "create", "entity_stake", stake.id, {
        "holder_entity_id": entity_id,
        "fund_entity_id": body.fund_entity_id,
    })
    session.commit()
    session.refresh(stake)
    return _stake_response(stake, {fund.id: fund})


@router.delete("/{entity_id}/stakes/{stake_id}")
def delete_stake(
    entity_id: int,
    stake_id: int,
    user: User = Depends(require_writer),
    session: Session = Depends(get_session),
) -> dict[str, str]:
    stake = session.get(EntityStake, stake_id)
    if stake is None or stake.holder_entity_id != entity_id:
        raise HTTPException(status_code=404, detail="Stake not found")
    record_audit(session, user.id, "delete", "entity_stake", stake_id, {
        "holder_entity_id": entity_id,
        "fund_entity_id": stake.fund_entity_id,
    })
    session.delete(stake)
    session.commit()
    return {"status": "deleted"}