# 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
from typing import Any, Optional
from unittest.mock import ANY

import pytest
import voluptuous
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent
from voluptuous import Schema

from czc.api.web.api import web_api
from czc.handler_context import HandlerContext, HandlerContextWithEvent
from czt.unittest import Context

from tests.sample_data import NAMESPACE, get_programmatic_api_test_event, get_test_event


@pytest.fixture(scope='function')
def context(mocker, environ):
    from czc import handler_context

    with Context(
        mocker,
        handler_context,
        ['event_bridge', 'aws_environment', 's3', 'feature_data', 'dyn'],
        extra={'handler_context': handler_context},
    ) as context:
        yield context


@pytest.mark.unit
@pytest.mark.parametrize('event', [{'foo': 'bar'}, get_test_event(), APIGatewayProxyEvent(get_test_event())])
def test_with_handler_context_should_handle_different_event_types(context, lambda_context, event):
    result_hc: Optional[HandlerContextWithEvent[Any]] = None

    @context.handler_context.with_handler_context()
    def handler(handler_context: HandlerContextWithEvent[Any]):
        nonlocal result_hc
        result_hc = handler_context
        return {}

    handler(event, lambda_context)
    assert result_hc == HandlerContextWithEvent(
        _organization_id=ANY,
        lambda_event=event,
        lambda_context=lambda_context,
        namespace=NAMESPACE,
        account_info=context.mock_aws_environment.resolve_account_info(),
        feature_conn=context.mock_feature_data.connect(),
        bucket=context.mock_s3.get_bucket_name(),
        dyn_table=context.mock_dyn.table(context.mock_dyn.get_table_name()),
        event_bus=context.mock_event_bridge.get_event_bus_name(),
    )


@pytest.mark.unit
@pytest.mark.parametrize(
    'event,expected_organization_id',
    [
        ({'organization_id': 'aaa-bbb-ccc'}, 'aaa-bbb-ccc'),
        ({'cz_organization_id': 'xxx-yyy-zzz'}, 'xxx-yyy-zzz'),
    ],
)
def test_with_handler_context_should_extract_organization_id_from_event(
    context, lambda_context, event, expected_organization_id
):
    result_hc: Optional[HandlerContext] = None

    @context.handler_context.with_handler_context(
        event_params=Schema(
            {
                voluptuous.Optional('organization_id'): str,
                voluptuous.Optional('cz_organization_id'): str,
            }
        )
    )
    def handler(handler_context: HandlerContext, **_kwargs):
        nonlocal result_hc
        result_hc = handler_context
        return {}

    handler(event, lambda_context)
    assert result_hc.organization_id == expected_organization_id


@pytest.mark.unit
@pytest.mark.parametrize(
    'event,expected_organization_id',
    [
        (get_test_event(organization_id='aaa-bbb-ccc'), 'aaa-bbb-ccc'),
        (APIGatewayProxyEvent(get_test_event(organization_id='xxx-yyy-zzz')), 'xxx-yyy-zzz'),
    ],
)
def test_with_handler_context_should_extract_organization_id_from_web_handler(
    context, lambda_context, event, expected_organization_id
):
    result_hc: Optional[HandlerContext] = None

    @web_api(path_schema={'userId': str})
    @context.handler_context.with_handler_context()
    def handler(handler_context: HandlerContext, **_kwargs):
        nonlocal result_hc
        result_hc = handler_context
        return {}

    handler(event, lambda_context)
    assert result_hc.organization_id == expected_organization_id


@pytest.mark.unit
@pytest.mark.parametrize(
    'event,expected_organization_id',
    [
        (get_programmatic_api_test_event(cz_organization_id='aaa-bbb-ccc', path_parameters={'userId': 'foo'}), 'aaa-bbb-ccc'),
        (get_test_event(organization_id='111-222-333', path_parameters={'userId': 'foo'}), '111-222-333'),
    ],
)
def test_with_handler_context_should_extract_organization_id_from_authorizer(
    context, lambda_context, event, expected_organization_id
):
    result_hc: Optional[HandlerContext] = None

    @context.handler_context.with_handler_context()
    def handler(handler_context: HandlerContext, **_kwargs):
        nonlocal result_hc
        result_hc = handler_context
        return {}

    handler(event, lambda_context)
    assert result_hc.organization_id == expected_organization_id
