import marshmallow

from users.app import ows_client


class CursorSchema(marshmallow.Schema):
    """Schema for cursor in pagination."""

    cursor = marshmallow.fields.String(required=True, allow_none=True)
    shorthand = marshmallow.fields.String(required=True, allow_none=True)


class RoleSchema(marshmallow.Schema):
    """Schema for a role."""

    role = marshmallow.fields.String(required=True)


class TenantRolesSchema(marshmallow.Schema):
    """Schema for tenant roles."""

    tenant_uuid = marshmallow.fields.String(required=True)
    tenant_type = marshmallow.fields.String(required=True)
    roles = marshmallow.fields.List(marshmallow.fields.Nested(RoleSchema), required=True)


class TenantRolesResponseSchema(marshmallow.Schema):
    """Schema for tenant roles response."""

    tenants = marshmallow.fields.Dict(
        keys=marshmallow.fields.String(),
        values=marshmallow.fields.Nested(TenantRolesSchema, required=True),
        required=True,
    )
    cursor = marshmallow.fields.Nested(CursorSchema, required=True)
    errors = marshmallow.fields.Dict()


class OwsPdpError(Exception):
    """Custom exception for ows-pdp errors."""

    def __init__(self, status_code: int, message: str):
        super().__init__(message)
        self.status_code = status_code


def get_tenant_roles_by_identity(identity_id: str) -> dict[str, list[str]]:
    """Call ows-pdp to get tenant_roles for a given identity ID."""
    tenant_roles = {}
    make_call = True
    cursor = ''
    while make_call:
        # Will use the JWT from the current request context
        response = ows_client.get(
            service_name='ows-pdp',
            path=f'identity/{identity_id}/roles/?cursor={cursor}',
        )
        if response.status_code != 200:
            raise OwsPdpError(
                status_code=response.status_code,
                message=f'Failed to get roles for identity {identity_id}: {response.text}',
            )
        tenant_roles_response = TenantRolesResponseSchema().load(response.json())
        tenant_roles.update(tenant_roles_response['tenants'])
        if tenant_roles_response['cursor']['cursor']:
            cursor = tenant_roles_response['cursor']['cursor']
        else:
            make_call = False
    return tenant_roles