from tempfile import TemporaryFile


def extract_name(description):
    return description.split("||", maxsplit=1)[0]


def migrate(migrator, database, fake=False, **kwargs):
    """
    This migration extracts incident name from whole mod_security message
    Centos6 version of sqlite does not have instr(), using slow python-based
    way
    """
    incident = migrator.orm["incident"]

    # FIXME: after migrating to peewee 3 remove the try..except block
    def select_incidents():
        try:
            yield from (
                incident.select(incident.id, incident.description)
                .where(incident.plugin == "modsec")
                .where(incident.description.contains("||"))
                .tuples()
                .iterator()
            )
        except RuntimeError:
            return

    with TemporaryFile(mode="w+") as f:
        for id_, desc in select_incidents():
            f.write("{},{}\n".format(id_, extract_name(desc)))
        f.seek(0)
        with database.atomic():
            for line in f:
                id_, name = line.split(",", maxsplit=1)
                incident.update(name=name.strip()).where(
                    incident.id == int(id_)
                ).execute()


def rollback(migrator, database, fake=False, **kwargs):
    """Write your rollback migrations here."""
    pass