"""Replay test generation from trace data."""

from __future__ import annotations

import textwrap
from typing import TYPE_CHECKING

from .models import get_function_alias

if TYPE_CHECKING:
    from ._trace_models import TracedFunction


def build_traced_arguments_call(
    func: TracedFunction,
    max_run_count: int,
) -> str:
    """Build the get_traced_arguments() call string."""
    parts = [
        "get_traced_arguments(",
        "trace_file=trace_file_path, ",
        f'function_name="{func.function_name}", ',
        f'file_name=r"{func.file_name}", ',
    ]
    if func.class_name is not None:
        parts.append(f'class_name="{func.class_name}", ')
    parts.append(f"num_to_get={max_run_count})")
    return "".join(parts)


def build_test_alias(func: TracedFunction) -> str:
    """Build the test function name alias."""
    if func.class_name is None:
        return get_function_alias(func.module_name, func.function_name)
    return get_function_alias(
        func.module_name,
        func.class_name + "_" + func.function_name,
    )


def build_replay_test_body(
    func: TracedFunction,
    max_run_count: int,
) -> str:
    """Build the body of a single replay test function."""
    call = build_traced_arguments_call(func, max_run_count)
    lines = [f"for arg_val_pkl in {call}:"]
    if func.class_name is None:
        alias = get_function_alias(func.module_name, func.function_name)
        lines.append("    args = pickle.loads(arg_val_pkl)")
        lines.append(f"    ret = {alias}(**args)")
    else:
        class_alias = get_function_alias(func.module_name, func.class_name)
        filter_line = ""
        if func.method_type == "classmethod":
            filter_line = '\n    args.pop("cls", None)'
        elif func.function_name == "__init__":
            filter_line = '\n    args.pop("__class__", None)'
        lines.append("    args = pickle.loads(arg_val_pkl)" + filter_line)
        method_name = (
            "." + func.function_name
            if func.function_name != "__init__"
            else ""
        )
        lines.append(f"    ret = {class_alias}{method_name}(**args)")
    lines.append("")
    return "\n".join(lines)


def create_trace_replay_test(
    trace_file: str,
    functions: list[TracedFunction],
    max_run_count: int = 100,
) -> str:
    """Generate a replay test file from a trace database."""
    imports = (
        "import pickle\n"
        "from codeflash_python.benchmarking._trace_db "
        "import get_traced_arguments\n"
    )
    function_imports: list[str] = []
    for function in functions:
        if not function.is_top_level:
            continue
        if function.class_name:
            alias = get_function_alias(
                function.module_name,
                function.class_name,
            )
            function_imports.append(
                f"from {function.module_name} import "
                f"{function.class_name} as {alias}"
            )
        else:
            alias = get_function_alias(
                function.module_name,
                function.function_name,
            )
            function_imports.append(
                f"from {function.module_name} import "
                f"{function.function_name} as {alias}"
            )
    imports += "\n".join(function_imports)
    functions_to_optimize = [
        f.function_name
        for f in functions
        if f.function_name != "__init__" and f.is_top_level
    ]
    metadata = (
        f"functions = {functions_to_optimize}\n"
        f'trace_file_path = r"{trace_file}"\n'
    )

    test_template = ""
    for func in functions:
        if not func.is_top_level:
            continue
        test_body = build_replay_test_body(func, max_run_count)
        alias = build_test_alias(func)
        formatted_test_body = textwrap.indent(test_body, "    ")
        test_template += f"def test_{alias}():\n{formatted_test_body}\n"

    return imports + "\n" + metadata + "\n" + test_template
