"""Convert behavioral instrumented tests to performance instrumented tests.

Both test types share the same test logic but differ in:
- Behavioral: captures (args, kwargs, return_value) to binary file for equivalence checking
- Performance: outputs timing via stdout markers for benchmarking, minimal overhead

This converter:
1. Replaces the codeflash_wrap function (binary file → stdout timing)
2. Removes inspect.signature().bind() overhead from test bodies
3. Strips unnecessary imports (inspect, dill, Path)
"""

from __future__ import annotations

import re

# The performance version of codeflash_wrap — outputs duration in stdout, no file I/O
PERF_WRAPPER = '''def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: Optional[str],
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    test_id = f"{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}"
    if not hasattr(codeflash_wrap, "index"):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f"{codeflash_line_id}_{codeflash_test_index}"
    exception = None
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(f"!$######{test_stdout_tag}######$!")
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f"!######{test_stdout_tag}:{codeflash_duration}######!")
    if exception:
        raise exception
    return return_value
'''


def convert_behavioral_to_perf(behavioral_source: str) -> str:
    """Convert a behavioral instrumented test file to a performance instrumented test file."""
    lines = behavioral_source.splitlines(keepends=True)

    # Step 1: Find where codeflash_wrap function starts and ends
    wrap_start = None
    wrap_end = None
    for i, line in enumerate(lines):
        if line.strip().startswith("def codeflash_wrap("):
            wrap_start = i
        # Find the end of codeflash_wrap — the next top-level def/class
        if wrap_start is not None and i > wrap_start + 5:
            stripped = line.lstrip()
            if (stripped.startswith("def ") or stripped.startswith("class ")) and not line[0].isspace():
                # But only if it's NOT the codeflash_wrap's own "def" line
                if "codeflash_wrap" not in stripped:
                    wrap_end = i
                    break
            # Also check for pytest marker at top level
            if stripped.startswith("@pytest") and not line[0].isspace():
                wrap_end = i
                break

    if wrap_start is None:
        # No codeflash_wrap found — can't convert
        return behavioral_source

    # If we didn't find the end, check for test functions
    if wrap_end is None:
        for i in range(wrap_start + 5, len(lines)):
            stripped = lines[i].lstrip()
            indent = len(lines[i]) - len(lines[i].lstrip())
            if indent == 0 and (stripped.startswith("def test_") or stripped.startswith("class Test")):
                wrap_end = i
                break
        if wrap_end is None:
            wrap_end = len(lines)

    # Step 2: Find the imports section (everything before codeflash_wrap)
    import_end = wrap_start
    # Walk backwards to find where imports end
    for i in range(wrap_start - 1, -1, -1):
        stripped = lines[i].strip()
        if stripped and not stripped.startswith("#"):
            import_end = i + 1
            break

    # Step 3: Build new file
    # Keep the project imports (from inference.xxx import ...) but remove behavioral-specific ones
    new_lines = []
    for i in range(import_end):
        line = lines[i]
        stripped = line.strip()
        # Skip behavioral-only imports
        if stripped in (
            "import inspect",
            "import dill as pickle",
            "from pathlib import Path",
            "import dill",
        ):
            continue
        # Skip codeflash_run_tmp_dir references
        if "codeflash_run_tmp_dir" in stripped:
            continue
        new_lines.append(line)

    # Add the perf wrapper
    new_lines.append("\n")
    new_lines.append(PERF_WRAPPER)
    new_lines.append("\n")

    # Step 4: Process test function bodies
    # Extract original args from inspect.signature().bind(args...) and pass them directly.
    # Two variants exist in the wild:
    #   v1: *_call__bound__arguments.args, **_call__bound__arguments.kwargs
    #   v2: **_call__bound__arguments.arguments
    pending_bind_args: str | None = None
    collecting_multiline_bind = False
    bind_lines: list[str] = []

    for i in range(wrap_end, len(lines)):
        line = lines[i]

        # Detect start of .bind( — may be single or multi-line
        if "_call__bound__arguments = inspect.signature(" in line:
            # Check if .bind() closes on this line
            match = re.search(r"\.bind\((.*)\)\s*$", line.strip())
            if match:
                pending_bind_args = match.group(1).strip()
            else:
                # Multi-line .bind( — start collecting
                collecting_multiline_bind = True
                bind_lines = [line]
            continue

        # Continue collecting multi-line .bind()
        if collecting_multiline_bind:
            bind_lines.append(line)
            joined = "".join(l.strip() for l in bind_lines)
            if ")" in line:
                match = re.search(r"\.bind\((.*)\)", joined)
                pending_bind_args = match.group(1).strip() if match else ""
                collecting_multiline_bind = False
                bind_lines = []
            continue

        # Skip apply_defaults line
        if "_call__bound__arguments.apply_defaults()" in line:
            continue

        # Replace v1: *_call__bound__arguments.args, **_call__bound__arguments.kwargs
        if "*_call__bound__arguments.args" in line:
            if pending_bind_args:
                line = line.replace(
                    "*_call__bound__arguments.args, **_call__bound__arguments.kwargs",
                    pending_bind_args,
                )
                line = line.replace(
                    ", *_call__bound__arguments.args, **_call__bound__arguments.kwargs",
                    f", {pending_bind_args}" if pending_bind_args else "",
                )
            else:
                line = line.replace(", *_call__bound__arguments.args, **_call__bound__arguments.kwargs", "")
                line = line.replace("*_call__bound__arguments.args, **_call__bound__arguments.kwargs", "")
            pending_bind_args = None
            new_lines.append(line)
            continue

        # Replace v2: **_call__bound__arguments.arguments
        if "**_call__bound__arguments.arguments" in line:
            if pending_bind_args:
                line = line.replace("**_call__bound__arguments.arguments", pending_bind_args)
                # Clean up trailing comma if args were removed
                line = re.sub(r",\s*,", ",", line)
                line = re.sub(r",\s*\)", ")", line)
            else:
                # No args — remove the line entirely if it's just the unpack
                stripped = line.strip()
                if stripped == "**_call__bound__arguments.arguments," or stripped == "**_call__bound__arguments.arguments":
                    continue
                line = line.replace(", **_call__bound__arguments.arguments", "")
                line = line.replace("**_call__bound__arguments.arguments,", "")
                line = line.replace("**_call__bound__arguments.arguments", "")
            pending_bind_args = None
            new_lines.append(line)
            continue

        new_lines.append(line)

    return "".join(new_lines)


def convert_file(behavioral_path: str, output_path: str) -> None:
    """Convert a behavioral test file to a performance test file."""
    from pathlib import Path

    source = Path(behavioral_path).read_text(encoding="utf-8")
    converted = convert_behavioral_to_perf(source)
    Path(output_path).write_text(converted, encoding="utf-8")


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} <behavioral_test.py> <output_perf_test.py>")
        sys.exit(1)
    convert_file(sys.argv[1], sys.argv[2])
