import asyncio
from pathlib import Path
import re
import aiofiles
from dbus_next import BusType
from dbus_next.aio import MessageBus
from inventory_handler.InventoryUtils import logger, RESCUE_CENTRAL_RAUCS, PERSISTENT_CENTRAL_RAUCS


async def readSections(filePath: Path) -> dict:
    """
    Read and parse sections from a RAUC-style INI file asynchronously.

    :param filePath: Path to the file to read.
    :type filePath: pathlib.Path
    :return: Dictionary where keys are section headers and values are section bodies.
    :rtype: dict
    """
    async with aiofiles.open(filePath, mode='r') as f:
        content = await f.read()
    sections = re.split(r"(?=\[.*?\])", content)
    parsed = {}
    for section in sections:
        if section.strip():
            lines = section.strip().splitlines()
            header = lines[0].strip()
            body = "\n".join(lines[1:])
            parsed[header] = body
    return parsed

async def writeSections(sections: dict, filePath: Path):
    """
    Write sections to a RAUC-style INI file asynchronously.

    :param sections: Dictionary of section headers and bodies to write.
    :type sections: dict
    :param filePath: Path to the file to write.
    :type filePath: pathlib.Path
    """
    content = "\n\n".join(f"{header}\n{body}" for header, body in sections.items()) + "\n"
    async with aiofiles.open(filePath, mode='w') as f:
        await f.write(content)

async def restartRAUCService():
    """
    Restart the RAUC service asynchronously via systemd D-Bus interface and verify that it
    has started properly by checking the 'Operation' property on the RAUC D-Bus interface.

    The function connects to the system bus, calls the systemd 'RestartUnit' method to restart
    'rauc.service', then waits up to 10 seconds for the RAUC service to re-register on D-Bus
    and report its operation state as 'idle' or another valid state.

    :raises TimeoutError: If RAUC does not come back online within the timeout period.
    """
    # Connect to the system bus
    bus = await MessageBus(bus_type=BusType.SYSTEM).connect()

    # Get a proxy object for systemd
    introspection = await bus.introspect('org.freedesktop.systemd1', '/org/freedesktop/systemd1')
    proxyObject = bus.get_proxy_object('org.freedesktop.systemd1',
                                        '/org/freedesktop/systemd1',
                                        introspection)

    manager = proxyObject.get_interface('org.freedesktop.systemd1.Manager')

    logger.info("Restarting RAUC service...")
    await manager.call_restart_unit('rauc.service', 'replace')

    RETRIES = 10
    DELAY = 1

    for attempt in range(RETRIES):
        try:
            raucObject = await bus.introspect('de.pengutronix.rauc', '/')
            raucProxy = bus.get_proxy_object('de.pengutronix.rauc', '/', raucObject)
            installer = raucProxy.get_interface('de.pengutronix.rauc.Installer')
            operation = await asyncio.wait_for(asyncio.shield(installer.get_operation()), timeout=2)
            logger.info(f"RAUC is up. Operation state: {operation}")
            return
        except Exception as e:
            if attempt < RETRIES - 1:
                logger.info(f"{attempt + 1}/{RETRIES}: RAUC is down ({e}). Retrying...")
                logger.info(f"Sleeping {DELAY}s before retry...")
                await asyncio.sleep(DELAY)
            else:
                logger.error("RAUC did not come back up in time.")
                raise TimeoutError("RAUC service restart timed out.") from e

async def syncRaucData(
        rescueCentralRaucsPath: Path = RESCUE_CENTRAL_RAUCS,
        persistentCentralRaucsPath: Path = PERSISTENT_CENTRAL_RAUCS
):
    """
    Synchronize RAUC section data between the rescue and persistent RAUC files asynchronously.

    If the rescue file is missing, the recovery section is removed from the persistent file.
    If both files exist, the recovery and rootfs sections are synchronized between them.

    :param rescueCentralRaucsPath: Path to the rescue RAUC file.
    :type rescueCentralRaucsPath: Path
    :param persistentCentralRaucsPath: Path to the persistent RAUC file.
    :type persistentCentralRaucsPath: Path
    """
    if not rescueCentralRaucsPath.exists():
        logger.info("Rescue file does not exist. Removing recovery block from persistent.")
        if not persistentCentralRaucsPath.exists():
            logger.warning("Persistent file not found. Nothing to do.")
            return

        sections = await readSections(persistentCentralRaucsPath)
        if "[slot.recovery.0]" in sections:
            del sections["[slot.recovery.0]"]
            await writeSections(sections, persistentCentralRaucsPath)
            logger.info("Removed [slot.recovery.0] from persistent.")
        else:
            logger.info("[slot.recovery.0] not found in persistent. Nothing removed.")
    else:
        logger.info("Rescue file exists. Syncing recovery block.")
        if not persistentCentralRaucsPath.exists():
            logger.warning("Persistent file not found. Cannot sync.")
            return

        rescueSections, persistentSections = await asyncio.gather(
            readSections(rescueCentralRaucsPath),
            readSections(persistentCentralRaucsPath)
        )

        if "[slot.recovery.0]" in rescueSections:
            persistentSections["[slot.recovery.0]"] = rescueSections["[slot.recovery.0]"]
            logger.info("Copied recovery block to persistent.")
        else:
            logger.warning("No recovery block found in rescue to copy.")

        await writeSections(persistentSections, persistentCentralRaucsPath)
        logger.info("Sync completed.")