#!/usr/bin/env python3
# Style normaliser for the port sources.
#   1. braces around single-statement if/else/for/while bodies
#   2. one statement per line (splits "a; b;" and "case X: stmt; break;")
#   3. snake_case identifiers -> camelCase (map built from the tree)
#   4. functions alphabetised (main last), static prototypes regenerated
import re, sys, os, glob

KEEP_SNAKE = set('''size_t ptrdiff_t int8_t int16_t int32_t int64_t uint8_t uint16_t uint32_t uint64_t
va_list wchar_t st_size st_mode off_t max_align_t int_fast8_t'''.split())

def camel(name):
    parts = name.split('_')
    out = parts[0]
    for p in parts[1:]:
        if not p:
            continue
        out += p[0].upper() + p[1:]
    return out

def strip_strings_comments(text):
    # replace string/char literal contents and comments with spaces (same length) for scanning
    out = []
    i = 0
    n = len(text)
    while i < n:
        c = text[i]
        if text.startswith('//', i):
            j = text.find('\n', i)
            if j < 0: j = n
            out.append(' ' * (j - i)); i = j
        elif text.startswith('/*', i):
            j = text.find('*/', i + 2)
            j = n if j < 0 else j + 2
            out.append(' ' * (j - i)); i = j
        elif c == '"' or c == "'":
            j = i + 1
            while j < n and text[j] != c:
                if text[j] == '\\': j += 1
                j += 1
            j = min(j + 1, n)
            out.append(c + ' ' * (j - i - 2) + c if j - i >= 2 else c); i = j
        else:
            out.append(c); i += 1
    return ''.join(out)

def build_rename_map(files):
    ids = {}
    strings = set()
    for f in files:
        t = open(f).read()
        code = strip_strings_comments(t)
        for m in re.finditer(r'\b[a-z][a-z0-9]*(?:_[a-z0-9]+)+\b', code):
            ids[m.group(0)] = ids.get(m.group(0), 0) + 1
        for s in re.findall(r'"(?:[^"\\]|\\.)*"', t):
            for m in re.finditer(r'\b[a-z][a-z0-9]*(?:_[a-z0-9]+)+\b', s):
                strings.add(m.group(0))
    mp = {}
    for k in ids:
        if k in KEEP_SNAKE or k.endswith('_t') or k.startswith('__'):
            continue
        mp[k] = camel(k)
    return mp

def apply_renames(text, mp):
    if not mp:
        return text
    pat = re.compile(r'\b(' + '|'.join(sorted(map(re.escape, mp), key=len, reverse=True)) + r')\b')
    # don't touch string literals or #include lines
    out = []
    for line in text.split('\n'):
        if line.lstrip().startswith('#include'):
            out.append(line); continue
        parts = re.split(r'("(?:[^"\\]|\\.)*")', line)
        for i in range(0, len(parts), 2):
            parts[i] = pat.sub(lambda m: mp[m.group(1)], parts[i])
        out.append(''.join(parts))
    return '\n'.join(out)

OP_END   = re.compile(r'(&&|\|\||<<|>>|[+\-*/|&=<>?:])\s*$')
OP_START = re.compile(r'^(&&|\|\||<<|>>|==|!=|<=|>=|\?|:|\+|-|\*|/|\||&|\.|->)')

def join_wrapped(text):
    """Join statements / declarations that were wrapped across lines."""
    lines = text.split('\n')
    out = []
    i = 0
    while i < len(lines):
        line = lines[i]
        s = line.strip()
        if s == '' or s.startswith('//') or s.startswith('#') or s.startswith('*') or s.startswith('/*') or s.endswith('\\'):
            out.append(line); i += 1; continue
        cur = line
        comments = []
        i += 1
        while i < len(lines):
            nostr = strip_strings_only(cur)
            if '//' in nostr:
                k = nostr.index('//')
                comments.append(cur[k + 2:].strip())
                cur = cur[:k].rstrip()
            code = strip_strings_comments(cur)
            nxt = lines[i]
            ns = nxt.strip()
            if ns == '' or ns.startswith('//') or ns.startswith('#') or ns.endswith('\\') or ns.startswith('/*'):
                break
            depth = code.count('(') - code.count(')') + code.count('[') - code.count(']')
            ends_op = OP_END.search(code) is not None and not code.rstrip().endswith('++') and not code.rstrip().endswith('--')
            starts_op = OP_START.match(ns) is not None and not (ns.startswith('-') and ns[1:2].isdigit() and code.rstrip().endswith(','))
            if code.rstrip().endswith(',') and depth <= 0:
                ends_op = False
            if code.rstrip().endswith('{') or code.rstrip().endswith('}') or code.rstrip().endswith(';'):
                if depth <= 0:
                    break
            if depth > 0 or ends_op or starts_op:
                cur = cur.rstrip() + ' ' + ns
                i += 1
                continue
            break
        if comments:
            cur = cur.rstrip() + '  // ' + '; '.join(c for c in comments if c)
        out.append(cur)
    return '\n'.join(out)

def strip_strings_only(text):
    out = []
    i = 0
    n = len(text)
    while i < n:
        c = text[i]
        if c == '"' or c == "'":
            j = i + 1
            while j < n and text[j] != c:
                if text[j] == '\\': j += 1
                j += 1
            j = min(j + 1, n)
            out.append(c + ' ' * (j - i - 2) + c if j - i >= 2 else c); i = j
        else:
            out.append(c); i += 1
    return ''.join(out)

DECL = re.compile(r'^(\s*)((?:static |const |unsigned |signed |struct |enum )*[A-Za-z_]\w*(?: \*+| )\**)([A-Za-z_]\w*(?:\[[^\]]*\])*)\s*(?:=\s*(.*?))?;\s*(//.*)?$')
ASSIGN = re.compile(r'^(\s*)([A-Za-z_][\w\.\[\]\(\)>-]*(?:\[[^\]]*\])*)\s*(=)\s*([^=].*);\s*(//.*)?$')

def align_runs(text):
    """Line up names in runs of declarations and '=' in runs of assignments."""
    lines = text.split('\n')
    out = []
    i = 0
    while i < len(lines):
        # declaration run
        m = DECL.match(lines[i])
        if m and not lines[i].strip().startswith('return') and '(' not in m.group(2):
            run = []
            j = i
            while j < len(lines):
                mm = DECL.match(lines[j])
                if not mm or mm.group(1) != m.group(1) or '(' in mm.group(2) or lines[j].strip().startswith('return'):
                    break
                run.append(mm); j += 1
            if len(run) >= 2:
                tw = max(len(r.group(2).rstrip()) for r in run)
                hasInit = [r for r in run if r.group(4) is not None]
                nw = max(len(r.group(3)) for r in hasInit) if hasInit else 0
                for r in run:
                    indent, typ, name, init, comment = r.groups()
                    typ = typ.rstrip()
                    # keep the '*' glued to the name
                    stars = ''
                    while typ.endswith('*'):
                        stars += '*'; typ = typ[:-1].rstrip()
                    line = indent + typ.ljust(tw - len(stars)) + ' ' + stars + name
                    if init is not None:
                        line = line.ljust(len(indent) + tw + 1 + nw) + ' = ' + init
                    line += ';'
                    if comment:
                        line += '  ' + comment
                    out.append(line)
                i = j
                continue
        # assignment run
        m = ASSIGN.match(lines[i])
        if m and not DECL.match(lines[i]):
            run = []
            j = i
            while j < len(lines):
                mm = ASSIGN.match(lines[j])
                if not mm or mm.group(1) != m.group(1) or DECL.match(lines[j]):
                    break
                run.append(mm); j += 1
            if len(run) >= 2:
                lw = max(len(r.group(2)) for r in run)
                for r in run:
                    indent, lhs, eq, rhs, comment = r.groups()
                    line = indent + lhs.ljust(lw) + ' = ' + rhs + ';'
                    if comment:
                        line += '  ' + comment
                    out.append(line)
                i = j
                continue
        out.append(lines[i]); i += 1
    return '\n'.join(out)

HEAD = re.compile(r'^(\s*)((?:\} else if|if|for|while) \((?:[^()]|\([^()]*\))*\)|\} else|else)\s+([^{};/][^;]*;)\s*(//.*)?$')

def fix_braces(text):
    lines = text.split('\n')
    out = []
    for line in lines:
        if line.rstrip().endswith('\\'):
            out.append(line); continue
        m = HEAD.match(line)
        if m and not line.rstrip().endswith('{') and 'for (' not in m.group(3):
            indent, head, stmt, comment = m.groups()
            out.append(f'{indent}{head} {{' + (f'  {comment}' if comment else ''))
            out.append(f'{indent}        {stmt.strip()}')
            out.append(f'{indent}}}')
            continue
        out.append(line)
    return '\n'.join(out)

CASE = re.compile(r'^(\s*)((?:case [^:]+|default):)\s+(\S.*;)\s*(//.*)?$')

def split_statements(text):
    lines = text.split('\n')
    out = []
    for line in lines:
        stripped = line.strip()
        if stripped.startswith('//') or stripped.startswith('#') or stripped.startswith('*') or stripped.endswith('\\'):
            out.append(line); continue
        m = CASE.match(line)
        if m:
            indent, label, rest, comment = m.groups()
            out.append(f'{indent}{label}' + (f'  {comment}' if comment else ''))
            for st in split_top(rest):
                out.append(f'{indent}        {st}')
            continue
        # generic "a; b;" lines (not for-headers, not braces-only)
        code = re.split(r'//', line, 1)
        body = code[0]
        comment = ('//' + code[1]) if len(code) > 1 else ''
        if 'for (' in body or '{' in body or '}' in body:
            out.append(line); continue
        sts = split_top(body.strip())
        if len(sts) > 1:
            indent = re.match(r'^(\s*)', line).group(1)
            for k, st in enumerate(sts):
                out.append(f'{indent}{st}' + (f'  {comment.strip()}' if (comment and k == len(sts) - 1) else ''))
            continue
        out.append(line)
    return '\n'.join(out)

def split_top(s):
    # split on ';' at paren depth 0, outside strings
    res = []; depth = 0; cur = ''; i = 0; instr = None
    while i < len(s):
        c = s[i]
        if instr:
            cur += c
            if c == '\\': cur += s[i+1]; i += 1
            elif c == instr: instr = None
        elif c in '"\'':
            instr = c; cur += c
        elif c in '([{': depth += 1; cur += c
        elif c in ')]}': depth -= 1; cur += c
        elif c == ';' and depth == 0:
            cur += c; res.append(cur.strip()); cur = ''
        else:
            cur += c
        i += 1
    if cur.strip():
        res.append(cur.strip())
    return res

FUNC = re.compile(r'^(?:static\s+)?(?:inline\s+)?(?:const\s+)?[A-Za-z_][\w\s\*]*?\b([A-Za-z_]\w*)\s*\(([^;]*)\)\s*\{\s*$')

def reorder_functions(text):
    lines = text.split('\n')
    n = len(lines)
    i = 0
    pre = []       # list of (kind, lines)
    funcs = []     # list of (name, lines, is_static, signature)
    buf = []       # pending non-function lines
    while i < n:
        line = lines[i]
        m = FUNC.match(line)
        if m and not line.startswith(' ') and not line.startswith('\t'):
            # pull preceding comment block (contiguous // lines) out of buf
            comment = []
            while buf and buf[-1].strip().startswith('//'):
                comment.insert(0, buf.pop())
            # trim trailing blank lines from buf
            while buf and buf[-1].strip() == '':
                buf.pop()
            if buf:
                pre.append(buf); buf = []
            # collect function body until a line that is exactly '}'
            body = [line]
            i += 1
            while i < n:
                body.append(lines[i])
                if lines[i] == '}':
                    i += 1
                    break
                i += 1
            sig = line.rstrip()[:-1].rstrip()  # drop '{'
            funcs.append((m.group(1), comment + body, line.startswith('static'), sig))
            continue
        buf.append(line); i += 1
    while buf and buf[-1].strip() == '':
        buf.pop()
    if buf:
        pre.append(buf)
    # remove existing static prototypes from preamble
    proto_re = re.compile(r'^static\s+[^=]*\b[A-Za-z_]\w*\s*\([^;]*\)\s*;\s*$')
    new_pre = []
    for chunk in pre:
        kept = [l for l in chunk if not proto_re.match(l)]
        while kept and kept[-1].strip() == '':
            kept.pop()
        while kept and kept[0].strip() == '':
            kept.pop(0)
        if kept:
            new_pre.append(kept)
    # regenerate prototypes for static functions
    protos = []
    for name, body, is_static, sig in funcs:
        if is_static:
            protos.append((name, sig))
    protos.sort(key=lambda t: t[0].lower())
    proto_lines = []
    if protos:
        # align names: split "static <type> name(args)"
        parsed = []
        for name, sig in protos:
            idx = sig.index(name + '(')
            rtype = sig[:idx].strip()
            args = sig[idx:]
            parsed.append((rtype, args))
        width = max(len(r) for r, a in parsed)
        for rtype, args in parsed:
            proto_lines.append(f'{rtype.ljust(width)} {args};')
    funcs.sort(key=lambda t: (t[0] == 'main', t[0].lower()))
    out = []
    for chunk in new_pre:
        out.extend(chunk); out.append(''); out.append('')
    if proto_lines:
        out.extend(proto_lines); out.append(''); out.append('')
    for k, (name, body, is_static, sig) in enumerate(funcs):
        out.extend(body)
        if k != len(funcs) - 1:
            out.append(''); out.append('')
    return '\n'.join(out).rstrip('\n') + '\n'

def process(path, mp, reorder=True):
    t = open(path).read()
    t = apply_renames(t, mp)
    if path.endswith('.h'):
        t = join_wrapped(t)
    if path.endswith('.c'):
        t = join_wrapped(t)
        t = fix_braces(t)
        t = split_statements(t)
        t = fix_braces(t)
        if reorder:
            t = reorder_functions(t)
        t = align_runs(t)
    open(path, 'w').write(t)

if __name__ == '__main__':
    root = sys.argv[1]
    files = sorted(glob.glob(root + '/src/*.c') + glob.glob(root + '/include/*.h') + glob.glob(root + '/tools/*.c'))
    mp = build_rename_map(files)
    if '--show-map' in sys.argv:
        for k in sorted(mp): print(k, '->', mp[k])
        sys.exit(0)
    noreorder = set(sys.argv[2:])
    for f in files:
        process(f, mp, reorder=os.path.basename(f) not in noreorder)
    print('processed', len(files), 'files;', len(mp), 'renames')