"""Memory-related fix executors."""

from __future__ import annotations

import xml.etree.ElementTree as ET
from typing import TYPE_CHECKING, Any

from mgf.vm.core.security.isolation.fixes._helpers import devices_element
from mgf.vm.core.security.isolation.types import isolation_fix

if TYPE_CHECKING:
    from mgf.vm.core.common.progress import ProgressReporter


__all__ = ["memory_add_balloon_none"]


@isolation_fix("memory.add_balloon_none")
def memory_add_balloon_none(
    *,
    target: ET.ElementTree,
    reporter: ProgressReporter,
    dry_run: bool = False,
    **_kw: Any,
) -> None:
    """Ensure ``<memballoon model='none'/>`` is present.

    If a balloon element exists with a different model, rewrite
    its ``model`` to ``"none"``. If no balloon element exists, add
    one. Idempotent.
    """
    devices = devices_element(target)
    if devices is None:
        reporter.report("memory.add_balloon_none: no <devices> element")
        return

    existing = devices.findall("memballoon")
    if not existing:
        if dry_run:
            reporter.report(
                "memory.add_balloon_none: (dry-run) would add <memballoon model='none'/>",
            )
        else:
            ET.SubElement(devices, "memballoon", attrib={"model": "none"})
            reporter.report("memory.add_balloon_none: added <memballoon model='none'/>")
        return

    changed = 0
    for balloon in existing:
        if balloon.get("model") == "none":
            continue
        if dry_run:
            reporter.report(
                f"memory.add_balloon_none: (dry-run) would set model='none' "
                f"on existing memballoon (was {balloon.get('model')!r})",
            )
        else:
            balloon.set("model", "none")
        changed += 1

    if changed == 0:
        reporter.report("memory.add_balloon_none: already satisfied")
        return
    if not dry_run:
        reporter.report(
            f"memory.add_balloon_none: set model='none' on {changed} memballoon(s)",
        )
