import peewee as pw


def migrate(migrator, database, fake=False, **kwargs):
    CountrySubnets = migrator.orm["country_subnets"]
    migrator.rename_field(CountrySubnets, "ip_net", "ip")
    migrator.add_fields(
        CountrySubnets,
        network_address=pw.IntegerField(null=True),
        netmask=pw.IntegerField(null=True),
        version=pw.IntegerField(null=True),
    )


def rollback(migrator, database, fake=False, **kwargs):
    CountrySubnets = migrator.orm["country_subnets"]
    migrator.rename_field(CountrySubnets, "ip", "ip_net")
    migrator.remove_fields(
        CountrySubnets, "network_address", "netmask", "version"
    )