# Copyright (c) 2022 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 Mapping

from toolz.dicttoolz import dissoc

import boto3
import cachetools.func
import pyfaaster.aws.handlers_decorators_v2 as faaster

from czc.constants import FEATURE_NAME
from czc.dynamodb import compose_key, dynamo_safe, get_expiry_timestamp
from czc.logging import get_logger

logger = get_logger()

SIMULATION_PREFIX = 'simulation'
TARGET_PREFIX = 'target'
ORDER_PREFIX = 'order'


@faaster.namespace_aware
def _get_table_name(NAMESPACE: str):
    return f'cz-{NAMESPACE}-{FEATURE_NAME}-data'


@cachetools.func.lru_cache(maxsize=1)
def _get_table():
    dynamodb = boto3.resource('dynamodb')
    return dynamodb.Table(_get_table_name())


def _to_public(simulation):
    return dissoc(simulation, 'PK', 'SK')


def _get_partition_key(cz_test_key):
    return compose_key(SIMULATION_PREFIX, cz_test_key)


def _get_sort_key(target: str, order: int):
    return compose_key(TARGET_PREFIX, target, ORDER_PREFIX, str(order))


def _get_primary_key(cz_test_key: str, target: str, order: int):
    return {
        'PK': _get_partition_key(cz_test_key),
        'SK': _get_sort_key(target, order)
    }


def get(test_key: str, target: str):
    if not test_key:
        return None

    table = _get_table()
    simulation_data = table.get_item(Key=_get_primary_key(test_key, target, 0)).get('Item')
    if simulation_data:
        logger.info(f'"Found simulation record. test_key={test_key} target={target} simulation={simulation_data}')

    return _to_public(simulation_data) if simulation_data else None


def put(test_key: str, target: str, data: Mapping):
    if not test_key:
        return None

    table = _get_table()

    record = {
        'ttl': dynamo_safe(get_expiry_timestamp()),  # Putting this first means it can be overridden with 'data'
        **data,
        **_get_primary_key(test_key, target, 0)
    }

    table.put_item(Item=record)
