#!/usr/bin/env python3
import argparse
import ast
from collections.abc import Iterable
from pathlib import Path


def _indent_lines(lines: Iterable[str], indent: int) -> str:
    prefix = ' ' * indent
    return '\n'.join(f'{prefix}{line}' if line else '' for line in lines)


def _as_doctest_expected(expr: ast.expr) -> str:
    try:
        value = ast.literal_eval(expr)
    except Exception:
        return ast.unparse(expr)
    return repr(value)


def _rewrite_pipe_to_data_last(expr: ast.expr) -> str | None:
    """Rewrite `R.pipe(data, R.fn())` into `R.fn()(data)` when possible."""
    if not isinstance(expr, ast.Call):
        return None
    if not isinstance(expr.func, ast.Attribute):
        return None
    if not isinstance(expr.func.value, ast.Name) or expr.func.value.id != 'R':
        return None
    if expr.func.attr != 'pipe':
        return None
    if len(expr.args) != 2:
        return None

    data_expr, fn_expr = expr.args

    # fn_expr must be `R.some_fn()` (i.e. Call(Attribute(Name('R'), ...), args=[]))
    if not isinstance(fn_expr, ast.Call):
        return None
    if not isinstance(fn_expr.func, ast.Attribute):
        return None
    if not isinstance(fn_expr.func.value, ast.Name) or fn_expr.func.value.id != 'R':
        return None
    if fn_expr.args or fn_expr.keywords:
        return None

    fn_name = fn_expr.func.attr
    return f"R.{fn_name}()({ast.unparse(data_expr)})"


def _extract_assert_eq(stmt: ast.stmt) -> tuple[ast.expr, ast.expr] | None:
    if not isinstance(stmt, ast.Assert):
        return None
    test = stmt.test
    if not isinstance(test, ast.Compare):
        return None
    if len(test.ops) != 1 or not isinstance(test.ops[0], ast.Eq):
        return None
    if len(test.comparators) != 1:
        return None
    return test.left, test.comparators[0]


def _group_for_method_name(name: str) -> str | None:
    if 'data_first' in name:
        return 'Data first:'
    if 'data_last' in name:
        return 'Data last:'
    return None


def convert_test_file_to_doctest(path: Path, indent: int) -> str:
    tree = ast.parse(path.read_text(), filename=str(path))

    groups: dict[str, list[tuple[str, str]]] = {'Data first:': [], 'Data last:': []}

    for node in tree.body:
        if not isinstance(node, ast.ClassDef):
            continue
        for fn in node.body:
            if not isinstance(fn, ast.FunctionDef):
                continue
            group = _group_for_method_name(fn.name)
            if group is None:
                continue

            for stmt in fn.body:
                extracted = _extract_assert_eq(stmt)
                if extracted is None:
                    continue
                actual_expr, expected_expr = extracted

                rewritten = _rewrite_pipe_to_data_last(actual_expr)
                actual_text = rewritten if rewritten is not None else ast.unparse(actual_expr)
                expected_text = _as_doctest_expected(expected_expr)
                groups[group].append((actual_text, expected_text))

    lines: list[str] = []
    for group_name in ('Data first:', 'Data last:'):
        if not groups[group_name]:
            continue
        if lines:
            lines.append('')
        lines.append(group_name)
        for actual_text, expected_text in groups[group_name]:
            lines.append(f">>> {actual_text}")
            lines.append(expected_text)

    return _indent_lines(lines, indent)


def main() -> None:
    parser = argparse.ArgumentParser(description='Convert remedapy unit tests into doctest-style examples.')
    parser.add_argument('target', help='Test file path or function name (e.g. "uncapitalise").')
    parser.add_argument('--indent', type=int, default=4)
    args = parser.parse_args()

    target_path = Path(args.target)
    if target_path.exists() and target_path.is_file():
        test_file = target_path
    else:
        name = args.target
        if name.startswith('R.'):
            name = name[2:]
        if name.startswith('test_'):
            name = name[5:]
        if name.endswith('.py'):
            name = name[:-3]
        repo_root = Path(__file__).resolve().parent.parent
        test_file = repo_root / 'src' / 'tests' / f'test_{name}.py'
        if not test_file.exists():
            raise FileNotFoundError(f'Could not find test file: {test_file}')

    print(convert_test_file_to_doctest(test_file, indent=args.indent))


if __name__ == '__main__':
    main()
