from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.modules.clients.association_models import ClientAssociation


def get_active_association(db: Session, client_id: int):
    stmt = (
        select(ClientAssociation)
        .where(ClientAssociation.client_id == client_id)
        .limit(1)
    )
    return db.execute(stmt).scalar_one_or_none()


def ensure_active_association(db: Session, client_id: int):
    row = get_active_association(db, client_id)
    if row:
        return row

    row = ClientAssociation(
        client_id=client_id,
        association_type="firm",
        created_source="system_admin",
    )
    db.add(row)
    db.commit()
    db.refresh(row)
    return row


def update_association_fields(db: Session, client_id: int, **fields):
    row = ensure_active_association(db, client_id)
    for key, value in fields.items():
        if hasattr(row, key):
            setattr(row, key, value)
    db.add(row)
    db.commit()
    db.refresh(row)
    return row