from app.application.exceptions import NotFoundError
from app.domain.entities.company import Company
from app.domain.entities.user_registration import UserRegistration
from app.domain.repositories.company_repository import CompanyRepository


async def resolve_company_for_user(
    repository: CompanyRepository,
    user: UserRegistration,
    company_id: str,
) -> Company:
    """Allow company owners and company_users members to access the same company."""
    company = await repository.get_for_user(company_id, user.id or "")
    if company:
        return company

    member_company_id = user.company_id
    if member_company_id:
        member_company = await repository.get_by_id(member_company_id)
        if member_company and (
            member_company.id == company_id or member_company.company_id == company_id
        ):
            return member_company

    raise NotFoundError("Company not found")