#!/usr/bin/env python3

import json
import re
import argparse
import os
import copy
import sys
import subprocess


VERSION = "2.2.0"


def warn(msg):
    sys.stderr.write(f"image2json: warning: {msg}\n")


IMAGE_KEYS = {"IGconf_device_class",
              "IGconf_device_variant",
              "IGconf_device_storage_type",
              "IGconf_device_sector_size",
              "IGconf_image_version",
              "IGconf_image_outputdir"}


IMAGE_KEY_TYPES = {
    "IGconf_device_sector_size": int,
    # add more as needed
}


top_template = {
    "IGversion": VERSION,
    "IGmeta": [],
    "attributes": {
        "image-name": "unknown",
        "image-size": 0,
        "image-palign-bytes": "unknown",
    },
    "layout": {
        "partitiontable": [],
        "partitionimages": [],
        "provisionmap": []
    }
}


partition_template =  {
    "name": "default",
    "in-partition-table": "true",
}


# mke2fs and others specify the UUID with -U
def get_uuid(extraargs):
    match = re.search(r"-U\s+([a-fA-F0-9\-]+)", extraargs)
    return match.group(1).lower() if match else None


# mkdosfs specifies volume ID with -i
def get_vfat_uuid(extraargs):
    match = re.search(r"-i\s+([a-fA-F0-9\-]+)", extraargs)
    if match:
        uuid = match.group(1).upper()
        return f"{uuid[:4]}-{uuid[4:]}"
    return None


# genimage key value mapping
def map_genimage(key, value):
    gmap = {
        # https://github.com/pengutronix/genimage#the-image-section
        "partition-type-uuid":{
            ("L", "linux"):       "0FC63DAF-8483-4772-8E79-3D69D8477DE4",
            ("S", "swap"):        "0657FD6D-A4AB-43C4-84E5-0933C84B4F4F",
            ("H", "home"):        "933AC7E1-2EB4-4F13-B844-0E14E2AEF915",
            ("U", "esp", "uefi"): "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
            ("R", "raid"):        "A19D880F-05FC-4D3B-A006-743F0F84911E",
            ("V", "lvm"):         "E6D6D379-F507-44C2-A23C-238F2A3DF928",
            ("F", "fat32"):       "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7"
        }
    }
    if key in gmap:
        translations = gmap[key]
        for possible_values, translated_value in translations.items():
            if value in possible_values:
                return translated_value

    return value


# General single line based key value reader with options
def read_kv(line, exclude_keys=None, translate_fn=None):
    if exclude_keys is None:
        exclude_keys = set() # Default - no exclusions

    for match in re.finditer(r"^\s*([\w-]+)\s*=\s*(\"[^\"]*\"|'[^']*'|[^#\n]+)", line, re.MULTILINE):
        key, value = match.groups()
        if key in exclude_keys:
            return (None, None)

        value = value.strip().strip('"').strip("'") # Remove quotes if present

        if translate_fn:
            value = translate_fn(key, value)

    return (key, value)


# Convert a libconfuse struct to nested dictionary with optional translate
# function and exclude filter for key value pairs
def confuse2dict(config_path, exclude_keys=None, tfn=None):
    with open(config_path, "r") as f:
        lines = f.readlines()

    root_structure = {}
    section_stack = [root_structure]

    # Match entry of section - type and optional subtype
    section_delimiter = r"([a-zA-Z0-9\-]+)(?:\s+([^\s]+))?\s*\{"

    for line in lines:
        line = line.strip()

        # remove any comments
        line = re.sub(r"\s*#.*", "", line).strip()
        if line == "" or not line:
            continue

        section = re.search(section_delimiter, line)
        if section:
            section_type = section.group(1)
            section_subtype = section.group(2)

            new_section = {}

            if section_subtype:
                if section_type not in section_stack[-1]:
                    section_stack[-1][section_type] = {} # new
                section_stack[-1][section_type][section_subtype] = new_section
                section_stack.append(new_section)
            else:
                section_stack[-1][section_type] = new_section
                section_stack.append(new_section)

            if '}' in line: # end of section on same line
                section_stack.pop()

            continue

        if "}" in line: # end of section
            if len(section_stack) > 1:
                section_stack.pop()
            else:
                raise ValueError("Unexpected closing brace")
            continue

        key, value = read_kv(line, exclude_keys, tfn)
        if key and value:
            if section_stack:
                section_stack[-1][key] = value

    if len(section_stack) != 1:
        raise ValueError("Parse error. Mismatched brace?")

    return root_structure


# Recursive find all sections by type
def find_sections(stype, data, sections=None):
    if sections is None:
        sections = {}

    for k, v in data.items():
        if isinstance(v, dict):
            if k == stype:
                for name, data in v.items():
                    sections[name] = data
            else:
                find_sections(stype, v, sections)
    return sections


# Find all images by type
def find_images(itype, data):
    timages = {}
    images = find_sections("image", data)
    for iname, iattr in images.items():
        if itype in iattr:
            timages[iname] = iattr
    return timages


# Returns the path of a build asset. If the asset is not an absolute path, use
# the env var as the directory of where to look.
def get_artefact_path(file, dir_evar):
    if not os.path.isabs(file):
        dir = os.getenv(dir_evar)
        if not dir:
            raise ValueError(f"Not an env var: {dir_evar}")
            return None

        if not os.path.isdir(dir):
            raise FileNotFoundError(f"Dir not found: {dir}")
            return None

        fpath = os.path.join(dir, file)
    else:
        fpath = file

    if not os.path.isfile(fpath):
        raise FileNotFoundError(f"File not found: {fpath}")
        return None

    return fpath


# Returns the size of a build artefact file.
def get_artefact_size(file):
    f = get_artefact_path(file, "IGconf_image_outputdir")
    if f:
        return os.stat(f).st_size
    return None


def get_fs_label(section):
    label = section.get("label")
    if label:
        return label
    return None


def process_extfs(section):
    args = section.get("extraargs", "")
    values = {
        "fs_uuid": get_uuid(args),
        "fs_label": get_fs_label(section)
    }
    return {k: v for k, v in values.items() if v is not None}


def process_vfat(section):
    args = section.get("extraargs", "")
    values = {
        "fs_uuid": get_vfat_uuid(args),
        "fs_label": get_fs_label(section)
    }
    return {k: v for k, v in values.items() if v is not None}


IMAGE_PROCESSORS = {
    "vfat":  process_vfat,
    "ext2":  process_extfs,
    "ext3":  process_extfs,
    "ext4":  process_extfs,
    "btrfs": process_extfs,
    "erofs": process_extfs
}


def parse_genimage_config(config_path):
    data = confuse2dict(config_path, {"exec-pre", "exec-post"}, map_genimage)

    # Anchor to hdimage as it's the root container of a disk image
    # https://github.com/pengutronix/genimage/?tab=readme-ov-file#hdimage
    hdimage_data = find_images("hdimage", data)
    if not hdimage_data:
        return (None,None,None)

    disk_attr = {}

    partitions = find_sections("partition", hdimage_data)
    images = find_sections("image", data)
    simgs = find_images("android-sparse", data)

    for img_name, img_data in hdimage_data.items():
        hdimage = img_data.get("hdimage")
        disk_attr["image-name"] = img_name
        disk_attr["image-size"] = hdimage.get("size") if hdimage.get("size") else 0
        disk_attr["image-palign"] = hdimage.get("align") if hdimage.get("align") else "0"
        sz = get_artefact_size(img_name)
        if sz:
            disk_attr["image-size"] = sz

    # Partitions declared outside the partition table are carried through to
    # the output, but as they have no table entry, they're excluded from the
    # association against the probed table.
    in_table = []
    for pname, pattr in partitions.items():
        if str(pattr.get("in-partition-table", "true")).lower() == "false":
            warn(f"Partition '{pname}' declared outside the partition table")
        else:
            in_table.append(pname)

    # Read the partition table from the disk image - partition sizes and GPT
    # attributes are derived from it.
    img_path = get_artefact_path(img_name, "IGconf_image_outputdir")
    try:
        res = subprocess.run(
                ["sfdisk", "--verify", "--json", img_path],
                capture_output=True,
                text=True,
                check=True
        )
    except Exception as e:
        raise RuntimeError(f"{e} Failed to read partition table from {img_path}")

    ptable = json.loads(res.stdout)
    pt_header = ptable.get("partitiontable", {})

    sfdisk_sectorsize = pt_header.get("sectorsize")
    if sfdisk_sectorsize is None:
        raise RuntimeError("sfdisk did not return sectorsize in partition table JSON")

    conf_sectorsize = os.environ.get("IGconf_device_sector_size")
    if conf_sectorsize and int(conf_sectorsize) != sfdisk_sectorsize:
        warn(f"IGconf_device_sector_size={conf_sectorsize} does not match "
             f"probed sectorsize={sfdisk_sectorsize}")

    # Sort sfdisk partitions by start offset for authoritative on-disk ordering
    sfdisk_parts = sorted(
        pt_header.get("partitions", []),
        key=lambda p: p.get("start", 0)
    )

    if pt_header.get("label", "").lower() == "dos" and len(sfdisk_parts) > 4:
        warn("MBR with more than 4 partitions requires an extended partition. "
             "GPT is recommended")

    # The association below is positional, so we must fail if our computed
    # 'in table' count differs to the probe.
    if len(in_table) != len(sfdisk_parts):
        raise ValueError(
            f"In partition table mismatch: genimage({len(in_table)}), "
            f"probe({len(sfdisk_parts)})"
        )

    # https://github.com/pengutronix/genimage?tab=readme-ov-file#the-image-configuration-options
    gtypes = ["android-sparse", "btrfs", "cpio", "cramfs", "ext2", "ext3",
              "ext4", "f2fs", "file", "FIT", "fip", "flash", "iso", "jffs2",
              "qemu", "rauc", "squashfs," "tar", "ubi", "vfat", "erofs"]

    is_gpt = pt_header.get("label", "").lower() == "gpt"

    # Associate all partitions with their images. Covers partitions in or out
    # of the ptable.
    for pname, pattr in partitions.items():
        # Initialise from template
        base = copy.deepcopy(partition_template)
        base["name"] = pname
        partitions[pname] = {**base, **pattr}

        piname = pattr.get("image")
        if piname not in images:
            continue

        # Found. Merge, config attributes win
        partitions[pname] = {**base, **images[piname], **pattr}

        # Tag image type
        for t in gtypes:
            s = find_sections(t, partitions[pname])
            if s:
                partitions[pname]["type"] = t
                # Invoke the processor for this type
                if t in IMAGE_PROCESSORS:
                    attr = IMAGE_PROCESSORS[t](s)
                    if attr:
                        partitions[pname].update(attr)

        # If this image has a sparse derivative, tag it
        for sname, sdata in simgs.items():
            simg = sdata.get("android-sparse")
            if piname == simg.get("image"):
                partitions[pname]["simage"] = sname

    # Attributes for partitions in the partition table come from the probe
    for pname, sfdisk_part in zip(in_table, sfdisk_parts):
        partitions[pname]["size"] = sfdisk_part["size"] * sfdisk_sectorsize

        # Populate GPT attributes from the partition table
        if is_gpt:
            gpt_name = sfdisk_part.get("name")
            if gpt_name:
                partitions[pname]["partition-label"] = gpt_name
            gpt_uuid = sfdisk_part.get("uuid")
            if gpt_uuid:
                partitions[pname]["partition-uuid"] = gpt_uuid

    # Attributes for partitions outside the partition table come from the image
    for pname in [n for n in partitions if n not in in_table]:
        piname = partitions[pname].get("image")
        sz = get_artefact_size(piname) if piname else None
        if sz:
            partitions[pname]["size"] = sz
        else:
            warn(f"Partition '{pname}' has no partition table entry and no "
                 "image to size it from")

    # Remove filesystem attributes from all partitions
    for part in partitions.values():
        fs_key = str(part.get("type", "")).lower()
        part.pop(fs_key, None)

    # Prune non-essential partition table info
    pt = ptable.get("partitiontable")
    if isinstance(pt, dict):
        for key in ("partitions", "device", "unit", "firstlba", "lastlba", "sectorsize"):
            pt.pop(key, None)

    return (disk_attr, partitions, ptable)


def get_env_vars(prefix=None):
    if prefix:
        return {key: value for key,value in os.environ.items() if key.startswith(prefix)}
    return dict(os.environ)


# Returns IG config vars we want to bake into the json
def get_image_meta():
    vars = get_env_vars()
    image_vars = {}
    for key, value in vars.items():
        if key in IMAGE_KEYS:
            if key in IMAGE_KEY_TYPES:
                try:
                    image_vars[key] = IMAGE_KEY_TYPES[key](value)
                except ValueError:
                    image_vars[key] = None  # or handle error as needed
            else:
                image_vars[key] = value
    return image_vars


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
            description='JSON layout generator')

    parser.add_argument("-g", "--genimage",
                        help="Path to genimage config file",
                        required=True)

    parser.add_argument("-m", "--provisionmap",
                        help="Path to the JSON provisioning map",
                        type=argparse.FileType('r'),
                        required=False)

    args = parser.parse_args()
    genimage_file = args.genimage;

    # Base info
    attributes, genimage_partitions, ptable = parse_genimage_config(genimage_file)
    partition_json = genimage_partitions

    top_template["IGmeta"] = get_image_meta()
    top_template["attributes"]["image-name"] = os.path.basename(attributes.get("image-name"))
    top_template["attributes"]["image-size"] = attributes.get("image-size")
    top_template["attributes"]["image-palign-bytes"] = attributes.get("image-palign")
    top_template["layout"]["partitiontable"] = ptable["partitiontable"]
    top_template["layout"]["partitionimages"] = partition_json
    if args.provisionmap:
        with args.provisionmap as f:
            top_template["layout"]["provisionmap"] = json.load(f)

    print(json.dumps(top_template, indent=4))
