#!/usr/bin/env python3
import json
import os
import subprocess
from pathlib import Path
from collections import defaultdict

# === Image Type Handlers ===
class ImageBuilder:
    def __init__(self, cfg):
        self.cfg = cfg

    def build(self):
        raise NotImplementedError("Image builder must implement build()")


class VFatImage(ImageBuilder):
    def build(self):
        img_path = Path(self.cfg["image"])
        size = self.cfg["size"]

        print(f"[VFAT] Creating {img_path} ({size} bytes), label={self.cfg.get('label','')}")
        with open(img_path, "wb") as f:
            f.truncate(size)
        subprocess.check_call([
            "mkfs.vfat",
            "-n", self.cfg.get("label", "NO_LABEL"),
            str(img_path)
        ])
        print(f"[VFAT] Done: {img_path}")


class HdImage(ImageBuilder):
    def build(self):
        img_path = Path(self.cfg["image"])
        partitions = self.cfg.get("partitions", {})

        total_size = 0
        for part in partitions.values():
            end_offset = part["offset"] + part.get("size", 0)
            total_size = max(total_size, end_offset)

        print(f"[HDIMAGE] Creating {img_path} ({total_size} bytes)")
        with open(img_path, "wb") as f:
            f.truncate(total_size)

        for name, part in partitions.items():
            file_path = Path(part["file"])
            if not file_path.exists():
                raise FileNotFoundError(f"Partition {name}: file {file_path} not found")
            with open(file_path, "rb") as src, open(img_path, "r+b") as dst:
                data = src.read()
                if "size" in part and len(data) > part["size"]:
                    raise ValueError(f"Partition {name} too large: {len(data)} > {part['size']}")
                dst.seek(part["offset"])
                dst.write(data)
                print(f"[HDIMAGE] Wrote {name} @0x{part['offset']:X} ({len(data)} bytes)")

        print(f"[HDIMAGE] Done: {img_path}")


# === Factory ===
IMAGE_BUILDERS = {
    "vfat": VFatImage,
    "hdimage": HdImage,
}

def build_image(cfg):
    img_type = cfg["image-type"]
    if img_type not in IMAGE_BUILDERS:
        raise ValueError(f"Unknown image type: {img_type}")
    IMAGE_BUILDERS[img_type](cfg).build()


# === Dependency resolution ===
def resolve_dependencies(configs):
    """
    Returns build order based on dependencies between images.
    """
    name_to_cfg = {cfg["image"]: cfg for cfg in configs}
    deps = defaultdict(set)

    for cfg in configs:
        img_name = cfg["image"]
        if cfg["image-type"] == "hdimage":
            for part in cfg.get("partitions", {}).values():
                if part["file"] in name_to_cfg:  # file is another image
                    deps[img_name].add(part["file"])

    # Topological sort
    visited = set()
    order = []

    def visit(img):
        if img in visited:
            return
        for dep in deps[img]:
            visit(dep)
        visited.add(img)
        order.append(img)

    for cfg in configs:
        visit(cfg["image"])

    return [name_to_cfg[name] for name in order]


# === Main ===
def main():
    import argparse
    parser = argparse.ArgumentParser(description="Dependency-aware Image Generator")
    parser.add_argument("config", help="JSON config file")
    args = parser.parse_args()

    with open(args.config) as f:
        configs = json.load(f)

    build_order = resolve_dependencies(configs)
    print("[INFO] Build order:", [cfg["image"] for cfg in build_order])

    for cfg in build_order:
        build_image(cfg)


if __name__ == "__main__":
    main()
