#!/usr/bin/env python3
"""Convert a simple .hex image to Xilinx .coe format.

Supported input:
- One or more hex tokens per line (e.g. "37" or "0x37")
- Optional comments after '#' or '//'

By default, each token is written as one .coe entry (8-bit style memory init).
Use --word-bytes > 1 to pack byte tokens into wider words.
"""

from __future__ import annotations

import argparse
from pathlib import Path


def parse_tokens(path: Path) -> list[int]:
    values: list[int] = []
    for raw in path.read_text().splitlines():
        line = raw.split("//", 1)[0].split("#", 1)[0].strip()
        if not line:
            continue
        for token in line.replace(",", " ").split():
            token = token.strip()
            if not token:
                continue
            if token.lower().startswith("0x"):
                token = token[2:]
            values.append(int(token, 16))
    return values


def pack_words(
    byte_values: list[int], word_bytes: int, little_endian: bool
) -> tuple[list[int], int]:
    if word_bytes <= 0:
        raise ValueError("word_bytes must be >= 1")
    if word_bytes == 1:
        return byte_values[:], 2

    words: list[int] = []
    width = word_bytes * 2
    for i in range(0, len(byte_values), word_bytes):
        chunk = byte_values[i : i + word_bytes]
        if len(chunk) < word_bytes:
            chunk = chunk + [0] * (word_bytes - len(chunk))

        word = 0
        if little_endian:
            for b_idx, b in enumerate(chunk):
                word |= (b & 0xFF) << (8 * b_idx)
        else:
            for b in chunk:
                word = (word << 8) | (b & 0xFF)
        words.append(word)

    return words, width


def main() -> None:
    parser = argparse.ArgumentParser(description="Convert .hex to Xilinx .coe")
    parser.add_argument("input_hex", type=Path, help="Input .hex file")
    parser.add_argument("output_coe", type=Path, help="Output .coe file")
    parser.add_argument(
        "--word-bytes",
        type=int,
        default=1,
        help="Bytes per output word (default: 1)",
    )
    parser.add_argument(
        "--little-endian",
        action="store_true",
        help="Pack bytes little-endian when --word-bytes > 1 (default)",
    )
    parser.add_argument(
        "--big-endian",
        action="store_true",
        help="Pack bytes big-endian when --word-bytes > 1",
    )
    parser.add_argument(
        "--radix",
        type=int,
        default=16,
        choices=[2, 10, 16],
        help="COE radix (default: 16)",
    )
    parser.add_argument(
        "--depth",
        type=int,
        default=0,
        help="Optional output depth (pads with zeros up to this many words)",
    )
    args = parser.parse_args()

    if args.little_endian and args.big_endian:
        raise SystemExit("Choose only one of --little-endian or --big-endian")

    little_endian = True
    if args.big_endian:
        little_endian = False

    byte_values = parse_tokens(args.input_hex)
    words, hex_digits = pack_words(byte_values, args.word_bytes, little_endian)

    if args.depth > 0 and args.depth < len(words):
        raise SystemExit(
            f"Requested --depth={args.depth} but image has {len(words)} words"
        )
    if args.depth > len(words):
        words.extend([0] * (args.depth - len(words)))

    if args.radix == 16:
        data = [f"{w:0{hex_digits}X}" for w in words]
    elif args.radix == 10:
        data = [str(w) for w in words]
    else:
        width_bits = args.word_bytes * 8
        data = [f"{w:0{width_bits}b}" for w in words]

    out_lines = [
        f"memory_initialization_radix={args.radix};",
        "memory_initialization_vector=",
    ]
    if data:
        out_lines.extend(
            [f"{v}," for v in data[:-1]] + [f"{data[-1]};"]
        )
    else:
        out_lines.append("0;")

    args.output_coe.write_text("\n".join(out_lines) + "\n")


if __name__ == "__main__":
    main()