from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session
import pytest
import itertools
import logging
from moto import mock_aws
import os
import base64

os.environ["MIND_CASTLE_LOCAL_ENCRYPTION_SECRET"] = base64.urlsafe_b64encode(
    os.urandom(16)
).decode()

from mind_castle.stores import retrieve_secret
from mind_castle.exceptions import RetrieveSecretException
from models import (
    Base,
    SimpleMemoryModel,
    SimpleLocalEncryptionModel,
)

logger = logging.getLogger(__name__)
logger.level = logging.DEBUG


# TEST_STORES = [SimpleMemoryModel, SimpleAWSSecretsManagerModel, SimpleAWSKMSModel, SimpleJsonModel, SimpleLocalEncryptionModel]
TEST_STORES = [SimpleMemoryModel, SimpleLocalEncryptionModel]
TEST_SECRET_VALUES = [
    {"secret_data_key": "secret_data_value"},
    ["secret_data_value_1", "secret_data_value_2"],
    "secret_string",
    123,
    4.56,
    "'",
    '"',
    '\\"',
    "\\'",
    "\\\\",
]
TEST_UPDATED_SECRET_VALUES = [
    {"secret_data_key": "secret_updated_value"},
    ["secret_updated_value_1", "secret_updated_value_2"],
    "secret_updated_string",
    456,
    7.89,
    '"',
    "'",
    "\\'",
    '\\"',
    "\\\\",
]


# Create an in-memory SQLite database
engine = create_engine("sqlite://", echo=True)
Base.metadata.create_all(engine)

session = Session(engine)

mock = mock_aws()
mock.start()


@pytest.mark.parametrize(
    "model, secret_data", itertools.product(TEST_STORES, TEST_SECRET_VALUES)
)
def test_create_secret_values(model, secret_data):
    test_object = model(data=secret_data)
    session.add(test_object)
    session.commit()
    session.expire_all()  # Get rid of cache

    item = session.get(model, test_object.id)
    assert item.data == secret_data


@pytest.mark.parametrize(
    "model, secret_data", itertools.product(TEST_STORES, TEST_SECRET_VALUES)
)
def test_retrieve_secret_values(model, secret_data):
    session.expire_all()  # Get rid of cache
    test_objects = session.scalars(select(model)).all()

    assert secret_data in [x.data for x in test_objects]
    item = next(x for x in test_objects if secret_data == x.data)
    assert isinstance(item.data, type(secret_data))


@pytest.mark.parametrize(
    "model, secret_data", itertools.product(TEST_STORES, TEST_SECRET_VALUES)
)
def test_update_secret_values(model, secret_data):
    test_objects = session.scalars(select(model)).all()
    item = next(x for x in test_objects if secret_data == x.data)
    secret_data_idx = TEST_SECRET_VALUES.index(secret_data)
    # key_before = item.data._self_secret_details["key"]

    # Update the value
    item.data = TEST_UPDATED_SECRET_VALUES[secret_data_idx]
    session.commit()
    session.expire_all()  # Get rid of cache

    test_object = session.get(model, item.id)
    assert test_object.data == TEST_UPDATED_SECRET_VALUES[secret_data_idx]
    # Decided to create new secrets on update for now, see comments in _track_secret_changes()
    # assert test_object.data._self_secret_details["key"] == key_before


@pytest.mark.parametrize(
    "model, secret_data", itertools.product(TEST_STORES, TEST_UPDATED_SECRET_VALUES)
)
def test_delete_secret_values(model, secret_data):
    test_objects = session.scalars(select(model)).all()
    item = next(x for x in test_objects if secret_data == x.data)

    session.delete(item)
    session.commit()
    session.expire_all()  # Get rid of cache

    test_objects = session.scalars(select(model)).all()
    with pytest.raises(StopIteration):
        item = next(x for x in test_objects if secret_data == x.data)

    # Local storage store types (KMS etc) wont throw an error here
    # with pytest.raises(RetrieveSecretException):
    try:
        retrieve_secret(item.data._self_secret_details)
    except RetrieveSecretException:
        pass
