from __future__ import annotations

import argparse
import ast
from pathlib import Path


MARKER = "    # PRESERVATION_METHODS\n"


def method_span(node: ast.FunctionDef | ast.AsyncFunctionDef) -> tuple[int, int]:
    start = min((decorator.lineno for decorator in node.decorator_list), default=node.lineno)
    if node.end_lineno is None:
        raise ValueError(f"Missing end position for {node.name}")
    return start - 1, node.end_lineno


def move_methods(
    source_path: Path,
    class_name: str,
    target_path: Path,
    method_names: list[str],
) -> None:
    source = source_path.read_text(encoding="utf-8")
    tree = ast.parse(source, filename=str(source_path))
    owner = next(
        (
            node
            for node in tree.body
            if isinstance(node, ast.ClassDef) and node.name == class_name
        ),
        None,
    )
    if owner is None:
        raise ValueError(f"Class not found: {class_name}")

    methods = {
        node.name: node
        for node in owner.body
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
    }
    missing = [name for name in method_names if name not in methods]
    if missing:
        raise ValueError(f"Methods not found in {class_name}: {', '.join(missing)}")

    lines = source.splitlines(keepends=True)
    ordered = sorted((methods[name] for name in method_names), key=lambda node: node.lineno)
    blocks = ["".join(lines[start:end]).rstrip() for start, end in map(method_span, ordered)]

    for start, end in sorted(map(method_span, ordered), reverse=True):
        del lines[start:end]
        while start < len(lines) - 1 and lines[start] == "\n" and lines[start + 1] == "\n":
            del lines[start]

    target = target_path.read_text(encoding="utf-8")
    if target.count(MARKER) != 1:
        raise ValueError(f"Target must contain exactly one method marker: {target_path}")
    target = target.replace(MARKER, "\n\n".join(blocks) + "\n")

    source_path.write_text("".join(lines), encoding="utf-8")
    target_path.write_text(target, encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Mechanically move class methods between modules")
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--class-name", required=True)
    parser.add_argument("--target", type=Path, required=True)
    parser.add_argument("methods", nargs="+")
    args = parser.parse_args()
    move_methods(args.source, args.class_name, args.target, args.methods)


if __name__ == "__main__":
    main()