#!/usr/bin/env python3
# libmisc/wrap-cc - Wrapper around GCC to enhance the preprocessor
#
# Copyright (C) 2025  Luke T. Shumaker <lukeshu@lukeshu.com>
# SPDX-License-Identifier: AGPL-3.0-or-later

import os
import subprocess
import sys
import typing


def scan_tuple(
    text: str, beg: int, on_part: typing.Callable[[str], None] | None = None
) -> int:
    assert text[beg] == "("
    pos = beg + 1
    arg_start = pos
    parens = 1
    instring = False
    while parens:
        c = text[pos]
        if instring:
            match c:
                case "\\":
                    pos += 1
                case '"':
                    instring = False
        else:
            match c:
                case "(":
                    parens += 1
                case ")":
                    parens -= 1
                    if on_part and parens == 0 and text[beg + 1 : pos].strip():
                        on_part(text[arg_start:pos])
                case ",":
                    if on_part and parens == 1:
                        on_part(text[arg_start:pos])
                        arg_start = pos + 1
                case '"':
                    instring = True
        pos += 1
    assert text[pos - 1] == ")"
    return pos - 1


def unquote(cstr: str) -> str:
    assert len(cstr) >= 2 and cstr[0] == '"' and cstr[-1] == '"'
    cstr = cstr[1:-1]
    out = ""
    while cstr:
        if cstr[0] == "\\":
            match cstr[1]:
                case "n":
                    out += "\n"
                    cstr = cstr[2:]
                case "\\":
                    out += "\\"
                    cstr = cstr[2:]
                case '"':
                    out += '"'
                    cstr = cstr[2:]
        else:
            out += cstr[0]
            cstr = cstr[1:]
    return out


def preprocess(all_args: list[str]) -> typing.NoReturn:
    # argparse #################################################################
    _args = all_args

    def shift(n: int) -> list[str]:
        nonlocal _args
        ret = _args[:n]
        _args = _args[n:]
        return ret

    arg0 = shift(1)[0]
    common_flags: list[str] = []
    output_flags: list[str] = []
    positional: list[str] = []
    while _args:
        if len(_args[0]) > 2 and _args[0][0] == "-" and _args[0][1] in "IDU":
            _args = [_args[0][:2], _args[0][2:], *_args[1:]]
        match _args[0]:
            # Mode
            case "-E" | "-quiet":
                common_flags += shift(1)
            case "-lang-asm":
                os.execvp(all_args[0], all_args)
            # Search path
            case "-I" | "-imultilib" | "-isystem":
                common_flags += shift(2)
            # Define/Undefine
            case "-D" | "-U":
                common_flags += shift(2)
            # Optimization
            case "-O0" | "-O1" | "-O2" | "-O3" | "-Os" | "-Ofast" | "-Og" | "-Oz":
                common_flags += shift(1)
            case "-g":
                common_flags += shift(1)
            # Output files
            case "-MD" | "-MF" | "-MT" | "-dumpbase" | "-dumpbase-ext":
                output_flags += shift(2)
            case "-o":
                output_flags += shift(2)
            # Other
            case _:
                if _args[0].startswith("-"):
                    if _args[0].startswith("-std="):
                        common_flags += shift(1)
                    elif _args[0].startswith("-m"):
                        common_flags += shift(1)
                    elif _args[0].startswith("-f"):
                        common_flags += shift(1)
                    elif _args[0].startswith("-W"):
                        common_flags += shift(1)
                    else:
                        raise ValueError(f"unknown flag: {_args!r}")
                else:
                    positional += shift(1)
    if len(positional) != 1:
        raise ValueError("expected 1 input file")
    infile = positional[0]

    # enhance ##################################################################

    common_flags += ["-D", "__LIBMISC_ENHANCED_CPP__"]

    text = subprocess.run(
        [arg0, *common_flags, infile],
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=sys.stderr,
        check=True,
        text=True,
    ).stdout

    macros: dict[str, str] = {}

    marker = "__xx__LM_DEFAPPEND__xx__"
    pos = 0
    while (marker_beg := text.find(marker, pos)) >= 0:
        args: list[str] = []

        def add_arg(arg: str) -> None:
            nonlocal args
            args.append(arg)

        beg_paren = marker_beg + len(marker)
        end_paren = scan_tuple(text, beg_paren, add_arg)

        before = text[:marker_beg]
        # old = text[marker_beg : end_paren + 1]
        after = text[end_paren + 1 :]

        assert len(args) == 2
        k = unquote(args[0].strip())
        v = unquote(args[1].strip())
        if k not in macros:
            macros[k] = v
        else:
            macros[k] += " " + v

        text = before + after
        pos = len(before)

    common_flags += ["-D", marker + "=LM_EAT"]
    for k, v in macros.items():
        common_flags += ["-D", k + "=" + v]

    # Run, for-real ############################################################
    os.execvp(arg0, [arg0, *common_flags, *output_flags, infile])


def main(all_args: list[str]) -> typing.NoReturn:
    if len(all_args) >= 2 and all_args[0].endswith("cc1") and all_args[1] == "-E":
        preprocess(all_args)
    else:
        os.execvp(all_args[0], all_args)


if __name__ == "__main__":
    main(sys.argv[1:])