import logging

from defence360agent.contracts.config import IConfig, LocalConfig

logger = logging.getLogger(__name__)


def _fix_config(config_file):
    try:
        config = config_file.config_to_dict(normalize=False)

        if "DOS" not in config:
            return

        config["DOS"].pop("timeout", None)
        config["DOS"].pop("max_connections", None)

        config_file.dict_to_config(config, overwrite=True, validate=False)
    except Exception:
        logger.exception("Failed to remove fields")


def migrate(
    migrator,
    database,
    fake=False,
    config_file: IConfig = LocalConfig(),
    **kwargs
):
    if fake:
        return

    _fix_config(config_file)


def rollback(migrator, database, fake=False, **kwargs):
    pass
