from typing import Annotated

from aws_lambda_powertools.event_handler.api_gateway import Router
from aws_lambda_powertools.event_handler.openapi.params import Query
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair

from boto3clients import dynamodb_client
from config import USER_TABLE

router = Router()
dyn = DynamoDBPersistenceLayer(USER_TABLE, dynamodb_client)


@router.get('/<user_id>/orgs')
def get_orgs(
    user_id: str,
    start_key: Annotated[str | None, Query] = None,
    limit: Annotated[int, Query(ge=25)] = 25,
):
    r = dyn.collection.query(
        # Post-migration (users): uncomment the following line
        # key=KeyPair(user_id, 'ORG#'),
        key=KeyPair(user_id, 'orgs#'),
        start_key=start_key,
        limit=limit,
    )

    preferred, items = None, []

    for x in r['items']:
        if 'DEFAULT' in x['sk']:
            preferred = x.get('org_id')
        else:
            # Post-migration (users): remove the following lines
            _, org_id = x['sk'].split('#')
            items.append(x | {'sk': f'ORG#{org_id}'})

    return r | {'items': items} | ({'preferred_org_id': preferred} if preferred else {})