# Copyright (c) 2021 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import re
import warnings
from collections.abc import Iterable

from aws_lambda_powertools.event_handler.exceptions import ServiceError as PowertoolsServiceError

from czc import json
from czc.api.v2.exceptions import ForbiddenError, ServiceError
from czc.logging import get_logger
from czc.util import deep_get

logger = get_logger()

#: Anchored regex pattern for CloudZero web origins. Pass as ``cors_origins=`` to avoid hand-rolling variants.
CLOUDZERO_CORS_ORIGINS: tuple[str, ...] = (r'^https://.*\.cloudzero\.com$',)


def http_response(default_error_message=None, cors_origins: Iterable[str] = None):
    """Lambda handler decorator that shapes responses, maps exceptions, and (optionally) handles CORS.

    When ``cors_origins`` is provided, this decorator:
      - rejects requests whose ``Origin`` header does not match any pattern with a 403 response, and
      - attaches ``Access-Control-Allow-Origin`` / ``Access-Control-Allow-Credentials`` headers to
        every response shape (success and caught exceptions).

    When ``cors_origins`` is empty (default), origin handling is skipped — preserving backward
    compatibility for callers that don't opt in.

    When ``cors_origins`` is set and the origin matches, the validated origin is also injected
    into the wrapped handler as ``kwargs['request_origin']``, preserving the contract that
    ``allow_origin_response`` had so downstream decorators (e.g. ones that read it for logging
    or connection metadata) continue to work without modification.

    Args:
        default_error_message: Fallback message for unhandled exceptions (default ``'API Error'``).
        cors_origins: Iterable of regex patterns matched against the request ``Origin`` header.
    """
    if cors_origins is None:
        cors_origins = ()
    origin_patterns = tuple(cors_origins)
    enforce_origin = bool(origin_patterns)

    def _origin_match(event):
        request_origin = deep_get(event, 'headers', 'origin', ignore_case=True)
        if not enforce_origin:
            return True, request_origin
        matched = bool(request_origin) and any(re.match(p, str(request_origin)) for p in origin_patterns)
        return matched, request_origin

    def _cors_headers(matched, request_origin):
        if not enforce_origin or not matched or not request_origin:
            return {}
        return {
            'Access-Control-Allow-Origin': request_origin,
            'Access-Control-Allow-Credentials': 'true',
        }

    def _log_by_severity(err, status_code):
        if status_code >= 500:
            logger.exception(err)
        else:
            logger.warning(err)

    def http_response_handler(handler):
        def handler_wrapper(event, context, **kwargs):
            origin_matched, request_origin = _origin_match(event)

            if enforce_origin and not origin_matched:
                logger.warning(f'Invalid request origin: {request_origin}')
                err = ForbiddenError('Unknown origin.')
                return {
                    'statusCode': err.status_code,
                    'body': err.body,
                    'headers': {},
                }

            cors = _cors_headers(origin_matched, request_origin)

            if enforce_origin:
                kwargs['request_origin'] = request_origin

            try:
                res = handler(event, context, **kwargs)
                if not isinstance(res, dict):
                    raise Exception(f'Unsupported return type {type(res)}; response must be dict.')
                return {
                    'headers': {**res.get('headers', {}), **cors},
                    'statusCode': res.get('statusCode', 200),
                    'body': json.dumps(res['body'], iterable_as_array=True) if 'body' in res else None,
                }
            except PowertoolsServiceError as err:
                _log_by_severity(err, err.status_code)
                return {
                    'statusCode': err.status_code,
                    'body': json.dumps(err.msg, iterable_as_array=True),
                    'headers': cors,
                }
            except ServiceError as err:
                _log_by_severity(err, err.status_code)
                return {
                    'statusCode': err.status_code,
                    'body': err.body,
                    'headers': cors,
                }
            except Exception as err:
                # Try and handle HTTPResponseException like objects
                if hasattr(err, 'statusCode') and hasattr(err, 'body'):
                    _log_by_severity(err, err.statusCode)
                    return {
                        'statusCode': err.statusCode,
                        'body': json.dumps(err.body, iterable_as_array=True),
                        'headers': cors,
                    }
                logger.exception(err)
                return {
                    'statusCode': 500,
                    'body': json.dumps({'error': {'message': default_error_message or 'API Error'}}, iterable_as_array=True),
                    'headers': cors,
                }

        return handler_wrapper

    return http_response_handler


def allow_origin_response(*origins):
    """Deprecated. Use ``http_response(cors_origins=[...])`` instead.

    This decorator is preserved as a backward-compatibility shim. It still rejects unknown
    origins by raising ``ForbiddenError`` and attaches CORS headers on success — but the
    safer path is to let ``http_response`` own both concerns so decorator order can't escape
    a raised ``ForbiddenError`` past the response-shaping layer (see CP-40686).
    """
    warnings.warn(
        'allow_origin_response is deprecated; use http_response(cors_origins=[...]) instead.',
        DeprecationWarning,
        stacklevel=2,
    )

    def allow_origin_handler(handler):
        def handler_wrapper(event, context, **kwargs):
            logger.debug(f'Checking origin for event: {event}')

            request_origin = deep_get(event, 'headers', 'origin', ignore_case=True)
            if not any(re.match(o, str(request_origin)) for o in origins):
                logger.warning(f'Invalid request origin: {request_origin}')
                raise ForbiddenError('Unknown origin.')

            kwargs['request_origin'] = request_origin
            response = handler(event, context, **kwargs)

            if not isinstance(response, dict):
                raise TypeError(f'Unsupported response type {type(response)}; response must be dict for *_response decorators.')

            current_headers = response.get('headers', {})
            cors_headers = {'Access-Control-Allow-Origin': request_origin, 'Access-Control-Allow-Credentials': 'true'}
            response['headers'] = {**current_headers, **cors_headers}
            return response

        return handler_wrapper

    return allow_origin_handler
