import asyncio
import argparse
import json
import logging
import os
import re
import shutil
import subprocess
import sys
import tarfile
from urllib.request import Request, urlopen
from urllib.error import HTTPError
from pathlib import Path
import time

from .screen import provideFirmwareUpdateScreen
from .infoenums import DeviceRole, OtaState
from .sdcard import SdCard
from .settings import (
    ROOT_FS_ARCHIVE_NAME,
    MOUNT_POINT_PERSISTENT_DATA,
    PERSISTENT_DATA_PARTITION,
    ROOTFS_0_PARTITION,
    ROOTFS_1_PARTITION,
    MOUNT_POINT_ROOTFS_0,
    MOUNT_POINT_ROOTFS_1,
    SD_CARD_BLOCKDEVICE_NAME,
    MOUNT_POINT_BOOT,
    UBOOT_PARTITION,
    UBOOT_CONFIGURATION_SCRIPTS,
    MTD_BOOTLOADER_PATH,
    UBOOT_FILE_NAME_S,
    UBOOT_FILE_NAME_DL,
    EVTEC_ECP_DATA,
    MOUNT_POINT_BACKUP,
    BACKED_UP_LOCATIONS,
    CONFIGURATION_FOLDER_MAPPING,
    OTA_RESCUE_LOGGER_NAME,
    ECP_OTA_BACKUP_FOLDER,
    BACKUP_FOLDER_MAPPING,
    OTA_MAIN_BOARD_API_PORT,
    OTA_MAIN_BOARD_END_POINT,
    ECP_OTA_DOWNLOAD_FOLDER,
    STATUS_FILE_NAME,
    UPDATE_PICTURE,
)

EXIT_CODE_FAILED_PARTITIONING = 1
EXIT_CODE_FAILED_FLASHING = 2

ECP_CONFIG = Path("/mnt/persistent-data/evtec-ecp-data/config/config.json")

otalogger = logging.getLogger(OTA_RESCUE_LOGGER_NAME)


def setUpLogging(otaWorkDirectory: Path):
    logFilePath = otaWorkDirectory / (OTA_RESCUE_LOGGER_NAME + ".log")

    otalogger.setLevel(logging.DEBUG)

    fHandler = logging.FileHandler(logFilePath)
    fHandler.setLevel(logging.DEBUG)

    sHandler = logging.StreamHandler(stream=sys.stdout)
    sHandler.setLevel(logging.DEBUG)

    formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
    fHandler.setFormatter(formatter)
    sHandler.setFormatter(formatter)

    otalogger.addHandler(fHandler)
    otalogger.addHandler(sHandler)

    # print header
    otalogger.info("")
    otalogger.info("".ljust(100, "-"))
    otalogger.info(" {}".format("OTA Rescue Flow"))
    otalogger.info("".ljust(99, "-"))


def readConfig() -> dict:
    """
    Read the config file and return it as a plain dictionary.
    """
    otalogger.info(f"Reading configuration from {ECP_CONFIG}")

    if not ECP_CONFIG.is_file():
        otalogger.warning("Config file not found.")
        return {}
    try:
        with open(ECP_CONFIG) as f:
            config = json.load(f)
    except json.JSONDecodeError as err:
        otalogger.warning(f"Failed to load {ECP_CONFIG.as_posix()}. Error: {err}")
        return {}
    except Exception as err:
        otalogger.warning(f"Failed to load {ECP_CONFIG.as_posix()}. Error: {err}")
        return {}

    return config


def buildIpAddressFromConfig(serial: int, role: DeviceRole) -> str:
    """
    Generate the IP address in the 10.2.x.x network from the serial number and the role.

    :param serial: serial number of the device
    :param role: role of the board configuration
    :return: IP address of the device with the specified serial number and role.
    """

    ipAddress = "10.2.{}.{}"
    lowByte = int(((serial - 1) % 100) + 1)
    highByte = int(((serial - 1) / 100) % 100)

    if role in [DeviceRole.CONTROLLER_REMOTE_DISPLAY, DeviceRole.CONTROLLER_LOCAL_HMI]:
        highByte = int(((serial - 1) / 100) % 100)
    elif role == DeviceRole.REMOTE_HMI_1:
        highByte = int(((serial - 1) / 100) % 100) + 100
    elif role == DeviceRole.REMOTE_HMI_2:
        highByte = int(((serial - 1) / 100) % 100) + 200

    return ipAddress.format(highByte, lowByte)


def isOtaMainBoard(config) -> bool:
    # config check

    if config.get("Device", None) and (role := config.get("Device").get("Role")):
        # check if the current board is the main board of the OTA process
        isOtaMainBoard = role in [
            DeviceRole.CONTROLLER_REMOTE_DISPLAY,
            DeviceRole.CONTROLLER_LOCAL_HMI,
            DeviceRole.POWER_SINK_CONTROL,
        ]
        return isOtaMainBoard
    else:
        return False


async def sendOtaNotification(
    destinationIp: str, sourceIp: str, notification: OtaState
) -> bool:
    """
    Notify the main board that the display board has successfully booted on ECP4
    """
    message = {"ip": sourceIp, "message": notification}
    payload = json.dumps(message).encode()
    req = Request(
        f"http://{destinationIp}:{OTA_MAIN_BOARD_API_PORT}/{OTA_MAIN_BOARD_END_POINT}",
        headers={
            "Content-Type": "application/json",
            "Content-Length": str(len(payload)),
        },
        data=payload,
        method="POST",
    )
    NB_RETRIES = 10
    RETRY_INTERVAL_SECONDS = 10
    success = False
    for i in range(NB_RETRIES):
        try:
            with urlopen(req) as response:
                data = response.read()
                success = response.status == 200
            otalogger.info(f"Server replied to {notification} with {data}")
        except (HTTPError, ConnectionRefusedError, Exception) as err:
            otalogger.error(f"Error occurred while sending {notification}: {err}")
        finally:
            if success:
                otalogger.info(f"Successfully sent {notification} to {destinationIp}")
                break
            else:
                otalogger.info(f"Failed to send {notification} to {destinationIp}.")
                await asyncio.sleep(RETRY_INTERVAL_SECONDS)
    return success


def extractTarXz(archivePath: Path, extractTo: Path = Path(".")) -> list[str]:
    """
    Extract *.tar.xz archive using `xz` command-line tool and `tarfile` module.

    :param archivePath: File path of the archive to be extracted.
    :param extractTo: File path to extract the archive to.
    :return: List of extracted files.
    """
    # The .xz is decompressed, now we have a .tar file
    tarFilePath = archivePath.with_suffix("")

    try:
        subprocess.check_call(["xz", "-d", "-k", archivePath])
    except subprocess.CalledProcessError as e:
        otalogger.error("Error decompressing .xz file: {!r}".format(e))
        os.remove(tarFilePath)
        return []

    fileNames = []
    try:
        with tarfile.open(tarFilePath, "r") as tar:
            fileNames = tar.getnames()
            tar.extractall(path=extractTo)
            otalogger.debug("Extracted to {}".format(extractTo))
    except tarfile.TarError as e:
        otalogger.error("Error extracting .tar file: {!r}".format(e))

    os.remove(tarFilePath)

    return fileNames


def gatherCPUCores() -> int:
    try:
        CAT = "/bin/cat"
        CPU_INFO = "/proc/cpuinfo"

        cpuInfoProcess = subprocess.Popen(
            [CAT, CPU_INFO], stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )

        stdout, stderr = cpuInfoProcess.communicate()

        if cpuInfoProcess.returncode != 0:
            otalogger.error(
                f"Command failed with return code: {cpuInfoProcess.returncode}"
            )
            otalogger.error(f"Error message:\n\t{stderr}")
            # return 1 to used u-boot-solo as default
            return 1

        pattern = r"[p,P]rocessor[ \t]*:[ \t]+(\w+)"
        decodedStdOut = stdout.decode()
        results = re.findall(pattern, decodedStdOut)
        if results:
            return len(results)
        else:
            # return 1 to used u-boot-solo as default
            return 1
    except Exception as e:
        otalogger.error(f"Error message:\n\t{e}")
        # return 1 to used u-boot-solo as default
        return 1


def flash_bootloader(mtdDevicePath: Path, bootloaderPath: Path):
    """
    Flash a new version of the U-Boot bootloader.
    """
    otalogger.info("Erasing bootloader flash.")

    flashEraseProcess = subprocess.Popen(
        ["flash_erase", mtdDevicePath.absolute(), "0", "0"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )

    stdout, stderr = flashEraseProcess.communicate()

    if flashEraseProcess.returncode != 0:
        otalogger.error(
            f"Command failed with return code: {flashEraseProcess.returncode}"
        )
        otalogger.error(f"Error message:\n\t{stderr}")
        # continue the flashing process, the next command might not fail, and the system
        # should be in a consistent state

    otalogger.info("Starting bootloader flashing process.")

    kobsProcess = subprocess.Popen(
        ["kobs-ng", "init", "-x", "-v", bootloaderPath.absolute()],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd="/tmp",
    )

    stdout, stderr = kobsProcess.communicate()

    if kobsProcess.returncode != 0:
        otalogger.error(f"Command failed with return code: {kobsProcess.returncode}")
        otalogger.error(f"Error message:\n\t{stderr}")
        return False

    return True


def readJsonContentFromFile(filePath):
    """
    Read json file from filePath and return it as dictionary.

    :param filePath: path to json file
    :return: file content as dictionary
    """
    content = {}
    try:
        with open(filePath, "r") as file:
            content = json.loads(file.read())
    except FileNotFoundError:
        otalogger.info("File {} not found for reading json content".format(filePath))
    except json.decoder.JSONDecodeError:
        otalogger.info("Failed to read file {} as JSON".format(filePath))
    except TypeError as e:
        otalogger.info("Failed to read {} with {}".format(filePath, e))
    return content


def addStatusEntry(workdir, message):
    filePath = os.path.join(workdir, STATUS_FILE_NAME)

    if os.path.exists(filePath) and os.path.getsize(filePath) > 0:
        try:
            data = readJsonContentFromFile(filePath)
            messages = data.get("status", [])

        except (json.JSONDecodeError, KeyError):
            messages = []
    else:
        messages = []

    messages.append(message)

    # Write back the updated messages to the file
    with open(filePath, "w") as file:
        json.dump({"status": messages}, file)

    time.sleep(5)


def showUpdatePicture():
    picturePath = Path("/opt/otarescue", "resources", UPDATE_PICTURE)
    showPictureProcess = subprocess.Popen(
        args=["dd", f"if={picturePath}", "of=/dev/fb0"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )

    stdout, stderr = showPictureProcess.communicate()

    if showPictureProcess.returncode != 0:
        otalogger.error(f"showUpdatePicture failed: {showPictureProcess.stderr}")
        return 1


async def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("-d", "--dir", type=str, required=True, help="OTA work dir")
    parser.add_argument(
        "-p",
        "--partition",
        action="store_true",
        help="Enables deleting and creating partitions.",
    )
    args = parser.parse_args()

    otaWorkDirectory = Path(args.dir)

    if not otaWorkDirectory.exists():
        os.makedirs(otaWorkDirectory, exist_ok=True)

    setUpLogging(otaWorkDirectory)

    # TODO: step [Flash U-Boot to NAND partition]
    addStatusEntry(otaWorkDirectory, "Install and configure ECP4 bootloader...")
    bootSourcePath = Path("/boot")

    ubootName = UBOOT_FILE_NAME_S
    otalogger.info("Flashing {} bootloader.".format(ubootName))

    if not flash_bootloader(Path(MTD_BOOTLOADER_PATH), bootSourcePath / ubootName):
        otalogger.error("Failed to flash bootloader.")
        return EXIT_CODE_FAILED_FLASHING

    config = readConfig()
    isMainBoard = isOtaMainBoard(config)
    ownIpAddress = "UNKNOWN"
    mainBoardIpAddress = "UNKNOWN"
    skipNotify = False

    try:
        serial = config.get("Device").get("Serial").get("Serial")
        role = config.get("Device").get("Role")
        ownIpAddress = buildIpAddressFromConfig(serial, role)
        mainBoardIpAddress = buildIpAddressFromConfig(
            serial, DeviceRole.CONTROLLER_REMOTE_DISPLAY
        )
    except Exception as err:
        otalogger.warning(f"Failed to parse config: {err}")
        skipNotify = True
    else:
        if not isMainBoard:
            showUpdatePicture()
            await sendOtaNotification(
                mainBoardIpAddress, ownIpAddress, OtaState.NOTIFY_BOOT_RESCUE
            )

    if isMainBoard:
        addStatusEntry(otaWorkDirectory, "Starting OTA Rescue...")
        provideFirmwareUpdateScreen(otaWorkDirectory, mainBoardIpAddress)
        time.sleep(15)

    sdcard = SdCard(SD_CARD_BLOCKDEVICE_NAME)

    if args.partition:
        # TODO: step [prepare ECP4 partition layout]
        addStatusEntry(otaWorkDirectory, "Prepare ECP4 partition layout...")
        # generate bootloader and rauc partitions
        if not sdcard.generatePartitions():
            return EXIT_CODE_FAILED_PARTITIONING

        # TODO: step [prepare for ECP4 installation]
        # mount partitions on prepared mount points
        bootPartitionMountPath = Path(MOUNT_POINT_BOOT)
        rootFs0PartitionMountPath = Path(MOUNT_POINT_ROOTFS_0)
        rootFs1PartitionMountPath = Path(MOUNT_POINT_ROOTFS_1)
        persistentDataPartitionMountPath = Path(MOUNT_POINT_PERSISTENT_DATA)

        addStatusEntry(otaWorkDirectory, "Mounting u-boot partition...")
        sdcard.mountPartition(UBOOT_PARTITION, bootPartitionMountPath)

        addStatusEntry(otaWorkDirectory, "Mounting rootfs.0 partition...")
        sdcard.mountPartition(ROOTFS_0_PARTITION, rootFs0PartitionMountPath)

        addStatusEntry(otaWorkDirectory, "Mounting rootfs.1 partition...")
        sdcard.mountPartition(ROOTFS_1_PARTITION, rootFs1PartitionMountPath)
        # Hint: no need to mount persistent-data, as it is mounted via /etc/fstab

        rootFsArchivePath = (
            persistentDataPartitionMountPath
            / ECP_OTA_DOWNLOAD_FOLDER
            / ROOT_FS_ARCHIVE_NAME
        )

        message = "Extracting rootfs to rootfs.0 partition..."
        addStatusEntry(otaWorkDirectory, message)
        otalogger.info(message)
        extractTarXz(rootFsArchivePath, extractTo=rootFs0PartitionMountPath)

        message = "Extracting rootfs to rootfs.1 partition..."
        addStatusEntry(otaWorkDirectory, message)
        otalogger.info(message)
        extractTarXz(rootFsArchivePath, extractTo=rootFs1PartitionMountPath)

        message = "Writing U-Boot scripts to boot partition..."
        addStatusEntry(otaWorkDirectory, message)
        otalogger.info(message)
        for file in UBOOT_CONFIGURATION_SCRIPTS:
            sourceFile = bootSourcePath / file
            shutil.copy(sourceFile, bootPartitionMountPath)

        # TODO: step [restore ECP4 configuration]
        addStatusEntry(otaWorkDirectory, "Restore ECP4 configuration...")
        otalogger.info("Copying configuration files from backup location.")
        ecpWorkDir = persistentDataPartitionMountPath / EVTEC_ECP_DATA

        ecpWorkDir.mkdir(parents=True, exist_ok=True)
        for folder in BACKED_UP_LOCATIONS:
            folderPath = ecpWorkDir / folder
            otalogger.debug(f"Creating folder `{folderPath}`")
            folderPath.mkdir(parents=True, exist_ok=True)

        for src, dest in BACKUP_FOLDER_MAPPING:
            # directly copy backed up to ECP4 that are not config related
            sourcePath = persistentDataPartitionMountPath / ECP_OTA_BACKUP_FOLDER / src
            if sourcePath.is_dir():
                destinationPath = ecpWorkDir / dest
                otalogger.debug(
                    f"Copying files from `{sourcePath}` to `{destinationPath}`"
                )
                shutil.copytree(sourcePath, destinationPath, dirs_exist_ok=True)
            else:
                otalogger.info(f"No directory {sourcePath}. Skip backup")

        if not isMainBoard and not skipNotify:
            await sendOtaNotification(
                mainBoardIpAddress, ownIpAddress, OtaState.NOTIFY_COMPLETE_RESCUE
            )

        addStatusEntry(otaWorkDirectory, "Start OTA in ECP4...")
        otalogger.info("Rebooting system after erasing the bootloader.")
        os.system("reboot")

    else:
        otalogger.info(
            "`--partition` argument not given, just showing the current SD card partition layout."
        )
        sdcard.printOverview()

    return 0


if __name__ == "__main__":
    return_code = asyncio.run(main())
    sys.exit(return_code)
