"""Linking a GP entity to its investor account (so Assets can pull real balances)."""

from datetime import date

from tests.conftest import make_user
from ten31portal.models import CapitalAccountStatement, Entity, EntityType, UserRole


def test_link_entity_to_investor(auth_client, session):
    inv = make_user(session, username="ten31llc", role=UserRole.investor, name="Ten31 LLC")
    gp = Entity(name="Ten31 LLC", type=EntityType.gp)
    session.add(gp)
    session.commit()
    session.refresh(gp)

    linked = auth_client.patch(f"/api/entities/{gp.id}", json={"linked_user_id": inv.id})
    assert linked.status_code == 200, linked.text
    assert linked.json()["linked_user_id"] == inv.id

    # Unlink.
    unlinked = auth_client.patch(f"/api/entities/{gp.id}", json={"linked_user_id": None})
    assert unlinked.status_code == 200
    assert unlinked.json()["linked_user_id"] is None


def test_link_rejects_non_investor(auth_client, session):
    staff = make_user(session, username="ops2", role=UserRole.operations)
    gp = Entity(name="Mgmt Co", type=EntityType.mgmt_co)
    session.add(gp)
    session.commit()
    session.refresh(gp)
    resp = auth_client.patch(f"/api/entities/{gp.id}", json={"linked_user_id": staff.id})
    assert resp.status_code == 400


def test_asset_balances_include_household(auth_client, session):
    """Balances under linked (household) names of the linked account are included."""
    primary = make_user(session, username="ten31llc", role=UserRole.investor, name="Ten31 LLC")
    secondary = make_user(
        session, username="ten31llc_trust", role=UserRole.investor,
        name="Ten31 LLC Trust", primary_account_id=primary.id,
    )
    f1 = Entity(name="LTPF I", type=EntityType.fund)
    f2 = Entity(name="LTPF II", type=EntityType.fund)
    gp = Entity(name="Ten31 LLC", type=EntityType.gp, linked_user_id=primary.id)
    session.add_all([f1, f2, gp])
    session.commit()
    for x in (f1, f2, gp):
        session.refresh(x)

    session.add(CapitalAccountStatement(
        entity_id=f1.id, investor_user_id=primary.id, as_of_date=date(2026, 3, 31),
        ending_balance_cents=600_000,
    ))
    # This balance sits under the linked secondary name, not the primary.
    session.add(CapitalAccountStatement(
        entity_id=f2.id, investor_user_id=secondary.id, as_of_date=date(2026, 3, 31),
        ending_balance_cents=400_000,
    ))
    session.commit()

    resp = auth_client.get(f"/api/entities/{gp.id}/asset-balances")
    assert resp.status_code == 200, resp.text
    body = resp.json()
    assert body["linked_name"] == "Ten31 LLC"
    assert {b["entity_id"] for b in body["balances"]} == {f1.id, f2.id}


def test_asset_balances_unlinked_is_empty(auth_client, session):
    gp = Entity(name="Mgmt", type=EntityType.mgmt_co)
    session.add(gp)
    session.commit()
    session.refresh(gp)
    resp = auth_client.get(f"/api/entities/{gp.id}/asset-balances")
    assert resp.status_code == 200
    assert resp.json()["balances"] == []