#!/usr/bin/env python3
"""refactor_order.py — reorder a go-firewall backend so it matches the Manager interface.

Ordering rules:
  1. Types and constants stay at the top (the preamble).
  2. The New<Type> constructor is placed immediately after the preamble.
  3. Remaining functions are ordered to match the `type Manager interface` list.
  4. Any helper must appear before the first function that calls it; dependencies are
     therefore hoisted in front of the exported method that needs them.

Doc comments attached to a function are moved together with the function.

Usage:
  scripts/refactor_order.py apf_linux.go                  # dry-run to stdout
  scripts/refactor_order.py apf_linux.go --apply          # rewrite in place
  scripts/refactor_order.py apf_linux.go --apply --gofmt  # rewrite and format
"""
import argparse
import os
import re
import subprocess
import sys

# Regex used to parse a column-0 `func` declaration line. Groups:
#   1: receiver variable (None for a plain function)
#   2: receiver type (None for a plain function)
#   3: function name
DECL_RE = re.compile(r"^func\s+(?:\((?:(\w+)\s+)?\*?(\w+)\)\s+)?([A-Za-z_]\w*)\s*\(")


def split_blocks(text):
    """Split source into (preamble, [block_dict, ...]).

    Each block begins with the doc comment for the function and ends at the
    start of the next function's doc comment. This keeps comments attached to
    the function they document, even when the file has been reordered before.

    Doc comments are detected by their first line: a column-0 `//` line whose
    leading word is the function name. The first function in the file is used
    to delimit the preamble.
    """
    lines = text.split("\n")
    func_lines = [i for i, line in enumerate(lines) if line.startswith("func ")]
    if not func_lines:
        return text, []

    # Parse all function declarations.
    funcs = []
    for idx, fl in enumerate(func_lines):
        m = DECL_RE.match(lines[fl])
        if not m:
            die(f"could not parse declaration: {lines[fl]}")
        assert m
        funcs.append(
            {
                "recv_var": m.group(1),
                "recv_type": m.group(2),
                "name": m.group(3),
                "func_line": fl,
                "body_end": func_lines[idx + 1]
                if idx + 1 < len(func_lines)
                else len(lines),
            }
        )
    func_names = {f["name"] for f in funcs}

    # Find all doc comment blocks in the file. A column-0 `//` sequence is split
    # into separate doc comments at lines that start with another function name.
    doc_re = re.compile(r"//\s+([a-zA-Z_][a-zA-Z0-9_]*)\b")
    doc_comments = {}  # name -> list of comment blocks (each block is a list of lines)
    i = 0
    while i < len(lines):
        if not lines[i].startswith("//"):
            i += 1
            continue
        seq = []
        while i < len(lines) and lines[i].startswith("//"):
            seq.append(lines[i])
            i += 1

        current = []
        for line in seq:
            mm = doc_re.match(line)
            if mm and mm.group(1) in func_names and current:
                # Line starts a new doc comment. Save the current one if valid.
                first = current[0]
                fm = doc_re.match(first)
                if fm and fm.group(1) in func_names:
                    doc_comments.setdefault(fm.group(1), []).append(current)
                current = [line]
            else:
                current.append(line)
        if current:
            first = current[0]
            fm = doc_re.match(first)
            if fm and fm.group(1) in func_names:
                doc_comments.setdefault(fm.group(1), []).append(current)

    # Deduplicate and choose the longest doc comment for each function.
    func_comment = {}
    for fname, blocks in doc_comments.items():
        seen = set()
        unique = []
        for block in blocks:
            key = "\n".join(block)
            if key not in seen:
                seen.add(key)
                unique.append(block)
        if unique:
            func_comment[fname] = max(unique, key=len)

    # Preamble is everything before the first function's doc comment or func line.
    first_func = funcs[0]
    if first_func["name"] in func_comment:
        comment = func_comment[first_func["name"]]
        # Find the earliest occurrence of this comment block before the function.
        preamble_end = first_func["func_line"]
        for pos in range(first_func["func_line"] - len(comment), -1, -1):
            if lines[pos : pos + len(comment)] == comment:
                preamble_end = pos
                break
    else:
        preamble_end = first_func["func_line"]
    preamble_lines = lines[:preamble_end]
    preamble = "\n".join(preamble_lines)
    if preamble_lines:
        preamble += "\n"

    # Build blocks. Body is from func_line to next_func_line, with the next
    # function's doc comment stripped off. The last function also has trailing
    # blank lines stripped so no orphan comments are left at the end of file.
    blocks = []
    for idx, f in enumerate(funcs):
        name = f["name"]
        is_last = idx == len(funcs) - 1
        body_end = f["body_end"]

        if is_last:
            # Skip trailing blanks so orphan comments after the last function are
            # reachable.
            while body_end > f["func_line"] and lines[body_end - 1] == "":
                body_end -= 1

        # Strip trailing // lines (the next function's doc comment, or orphan
        # comments after the last function).
        while body_end > f["func_line"] and lines[body_end - 1].startswith("//"):
            body_end -= 1

        if is_last:
            # Strip any remaining blanks after the orphan comments.
            while body_end > f["func_line"] and lines[body_end - 1] == "":
                body_end -= 1

        body_lines = lines[f["func_line"] : body_end]
        body_text = "\n".join(body_lines)
        # Preserve a separator blank line that sits immediately before the next
        # function's doc comment. Python's split/join drops the blank when it is
        # the last element of the slice, so add it back explicitly.
        if (
            body_lines
            and body_lines[-1] == ""
            and any(line != "" for line in lines[body_end:])
        ):
            body_text += "\n"
        if body_text and not body_text.endswith("\n"):
            body_text += "\n"

        comment_text = ""
        if name in func_comment:
            comment_text = "\n".join(func_comment[name]) + "\n"

        blocks.append(
            {
                "recv_var": f["recv_var"],
                "recv_type": f["recv_type"],
                "name": name,
                "text": comment_text + body_text,
            }
        )

    return preamble, blocks


def parse_manager_interface(interface_path):
    """Return the method names of `type Manager interface` in order."""
    if not os.path.isfile(interface_path):
        die(f"interface file not found: {interface_path}")
    text = open(interface_path).read()
    m = re.search(r"type Manager interface \{(.*?)\n\}", text, re.DOTALL)
    if not m:
        die(f"type Manager interface not found in {interface_path}")
    assert m
    body = m.group(1)
    # Strip interface comments so they do not interfere with method detection.
    body = re.sub(r"(?m)^\s*//.*$", "", body)
    methods = re.findall(r"^\s+([A-Z]\w+)\s*\(", body, re.MULTILINE)
    return methods


def detect_backend_type_and_receiver(blocks):
    """Detect the dominant backend receiver type and variable name."""
    recv_counts = {}
    type_counts = {}
    for block in blocks:
        rv = block["recv_var"]
        rt = block["recv_type"]
        if rt:
            recv_counts[rv] = recv_counts.get(rv, 0) + 1
            type_counts[rt] = type_counts.get(rt, 0) + 1
    if not type_counts:
        die("could not detect a backend receiver type")
    backend_type = max(type_counts, key=lambda k: type_counts[k])
    receiver = max(recv_counts, key=lambda k: recv_counts[k]) if recv_counts else None
    return backend_type, receiver


def build_call_graph(blocks, receiver):
    """Return a map from function name to the set of in-file functions it calls."""
    names = {block["name"] for block in blocks}
    graph = {name: set() for name in names}

    method_call_re = None
    if receiver:
        method_call_re = re.compile(rf"\b{re.escape(receiver)}\.([A-Za-z_]\w*)\(")

    for block in blocks:
        name = block["name"]
        text = block["text"]
        deps = graph[name]

        # Receiver method calls: f.name(...)
        if method_call_re:
            for mm in method_call_re.finditer(text):
                called = mm.group(1)
                if called != name and called in names:
                    deps.add(called)

        # Plain function calls: name(...)
        # Exclude calls preceded by a dot, since those are method calls on some
        # other value and are already captured above for the backend receiver.
        for mm in re.finditer(r"(?<!\.)\b([A-Za-z_]\w*)\(", text):
            called = mm.group(1)
            if called != name and called in names:
                deps.add(called)

    return graph


def reorder(blocks, manager_order, constructor):
    """Return blocks in the required order."""
    backend_type, receiver = detect_backend_type_and_receiver(blocks)
    graph = build_call_graph(blocks, receiver)
    name_to_block = {block["name"]: block for block in blocks}
    names = set(name_to_block.keys())

    if constructor not in names:
        die(f"constructor {constructor} not found in file")

    missing = [m for m in manager_order if m not in names]
    if missing:
        die(f"Manager interface methods missing from file: {missing}")

    order = []
    placed = set()

    def place(name):
        if name in placed:
            return
        if name not in name_to_block:
            return
        for dep in sorted(graph.get(name, set())):
            place(dep)
        order.append(name)
        placed.add(name)

    # 1. Constructor after preamble, before other functions.
    place(constructor)

    # 2. Manager interface methods in interface order.
    for method in manager_order:
        place(method)

    # 3. Any leftover functions (e.g. helper-only internals not reachable from
    #    the public surface) are appended in dependency order.
    for name in sorted(names):
        if name not in placed:
            place(name)

    return [name_to_block[name] for name in order]


def die(msg):
    print("error: " + msg, file=sys.stderr)
    sys.exit(1)


def main():
    ap = argparse.ArgumentParser(
        description="Reorder a go-firewall backend to match the Manager interface and dependency order."
    )
    ap.add_argument("file", help="backend source file (e.g. apf_linux.go)")
    ap.add_argument(
        "--interface",
        default="firewall.go",
        help="file containing the Manager interface (default: firewall.go)",
    )
    ap.add_argument(
        "--constructor",
        help="constructor name override (default: New<Type> auto-detected)",
    )
    ap.add_argument("--apply", action="store_true", help="rewrite the file in place")
    ap.add_argument("--gofmt", action="store_true", help="run gofmt after rewriting")
    args = ap.parse_args()

    if not os.path.isfile(args.file):
        die(f"no such file: {args.file}")

    text = open(args.file).read()
    preamble, blocks = split_blocks(text)
    if not blocks:
        die("no functions found in file")

    backend_type, _ = detect_backend_type_and_receiver(blocks)
    constructor = args.constructor or ("New" + backend_type)

    manager_order = parse_manager_interface(args.interface)
    ordered = reorder(blocks, manager_order, constructor)
    new_text = preamble + "".join(block["text"] for block in ordered)

    if not args.apply:
        print(new_text, end="")
        return

    open(args.file, "w").write(new_text)
    if args.gofmt:
        subprocess.run(["gofmt", "-w", args.file])
    print(f"reordered {args.file}")


if __name__ == "__main__":
    main()