from __future__ import annotations

import json 
from typing import TYPE_CHECKING

import jwt 

from ..config import CognitoAuthConfig
from ..exceptions import ConfigurationError
from ..config import DEFAULT_REGION

if TYPE_CHECKING:
    from jwt import PyJWKClient 

def build_jwt_client(config: CognitoAuthConfig) -> 'PyJWKClient':
    return jwt.PyJWKClient(
        config.jwk_url(),
    )

def load_config_from_ssm(
    parameter_name: str,
    region: str = DEFAULT_REGION,
) -> CognitoAuthConfig:
    try:
        import boto3 
    except ImportError as e:
        raise ConfigurationError(
            f'boto3 not installed: {e}',
            status_code=400,
            details={'package': 'boto3'},
        ) from e 

    client_kw: dict[str, str] = {}
    if region != DEFAULT_REGION:
        client_kw['region_name'] = region 

    ssm = boto3.client('ssm', **client_kw)
    try:
        response = ssm.get_parameter(Name=parameter_name, WithDecryption=True)
    except Exception as e:
        raise ConfigurationError(
            f'Failed to load config from SSM: {e}',
            status_code=500,
            details={'parameter': parameter_name},
        ) from e 

    raw = response['Parameter']['Value']
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        raise ConfigurationError(
            f'Invalid JSON in SSM parameter: {e}',
            status_code=500,
            details={'parameter': parameter_name},
        ) from e 

    try:
        return CognitoAuthConfig(
            user_pool_id=data.get('user_pool_id'),
            region=data.get('region'),
            app_client_id=data.get('app_client_id'),
            group_to_role=data.get('group_to_role', {}),
        )
    except KeyError as e:
        raise ConfigurationError(
            f'Missing required field in SSM parameter: {e}',
            status_code=500,
            details={'parameter': parameter_name},
        ) from e 
        
        

    