#!/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

# Matches a function declaration together with its immediately preceding
# column-0 `//` doc comments. Groups:
#   3: receiver variable (None for a plain function)
#   4: receiver type (None for a plain function)
#   5: function name
FUNC_RE = re.compile(
    r"(?m)^((?://.*\n)*)^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 comments immediately preceding a `func` line
    and ends at the start of the next block, so comments stay attached to their
    function.
    """
    lines = text.split("\n")
    func_lines = [i for i, line in enumerate(lines) if line.startswith("func ")]
    if not func_lines:
        return text, []

    decl_re = re.compile(r"^func\s+(?:\((\w+)\s+\*?(\w+)\)\s+)?([A-Za-z_]\w*)\s*\(")

    def comment_start(func_line):
        i = func_line - 1
        # Allow a single blank line between the previous block and the doc comment.
        if i >= 0 and lines[i].strip() == "":
            i -= 1
        while i >= 0 and lines[i].startswith("//"):
            i -= 1
        return i + 1

    comment_starts = [comment_start(fl) for fl in func_lines]
    first_start = comment_starts[0]
    preamble_lines = lines[:first_start]
    preamble = "\n".join(preamble_lines)
    if preamble_lines:
        preamble += "\n"

    blocks = []
    for idx, fl in enumerate(func_lines):
        end_line = func_lines[idx + 1] if idx + 1 < len(func_lines) else len(lines)
        block_lines = lines[comment_starts[idx] : end_line]
        block_text = "\n".join(block_lines)
        if block_text and not block_text.endswith("\n"):
            block_text += "\n"

        m = decl_re.match(lines[fl])
        if not m:
            die(f"could not parse declaration: {lines[fl]}")
        assert m
        blocks.append(
            {
                "recv_var": m.group(1),
                "recv_type": m.group(2),
                "name": m.group(3),
                "text": block_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()