"""Line profiler utilities.

Adapted from line_profiler written by Enthought, Inc. (BSD License).
See https://github.com/pyutils/line_profiler.

CST transformers for adding line_profiler decorators and imports to
Python source files, plus a high-level orchestration function.
"""

from __future__ import annotations

import logging
from collections import defaultdict
from pathlib import Path
from typing import TYPE_CHECKING

import libcst as cst

from ..runtime._codeflash_wrap_decorator import get_run_tmp_file

if TYPE_CHECKING:
    from collections.abc import Sequence

    from .._model import FunctionSource, FunctionToOptimize

log = logging.getLogger(__name__)


class LineProfilerDecoratorAdder(cst.CSTTransformer):
    """Adds a decorator to a function matched by qualified name."""

    def __init__(self, qualified_name: str, decorator_name: str) -> None:
        """Initialize with the target qualified name and decorator to add."""
        super().__init__()
        self.qualified_name_parts = qualified_name.split(".")
        self.decorator_name = decorator_name
        self.context_stack: list[str] = []

    def visit_ClassDef(self, node: cst.ClassDef) -> None:  # noqa: N802
        """Push class name onto the context stack."""
        self.context_stack.append(node.name.value)

    def leave_ClassDef(  # noqa: N802
        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
    ) -> cst.ClassDef:
        """Pop class name from the context stack."""
        self.context_stack.pop()
        return updated_node

    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:  # noqa: N802
        """Push function name onto the context stack."""
        self.context_stack.append(node.name.value)

    def leave_FunctionDef(  # noqa: N802
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef:
        """Add the decorator if the current context matches the target."""
        if self.context_stack == self.qualified_name_parts:
            has_decorator = any(
                self._is_target_decorator(decorator.decorator)
                for decorator in original_node.decorators
            )
            if not has_decorator:
                new_decorator = cst.Decorator(
                    decorator=cst.Name(value=self.decorator_name)
                )
                updated_decorators = [
                    new_decorator,
                    *list(updated_node.decorators),
                ]
                updated_node = updated_node.with_changes(
                    decorators=tuple(updated_decorators)
                )
        self.context_stack.pop()
        return updated_node

    def _is_target_decorator(self, decorator_node: cst.BaseExpression) -> bool:
        """Check if a decorator node matches the target decorator name."""
        if isinstance(decorator_node, cst.Name):
            return decorator_node.value == self.decorator_name
        if isinstance(decorator_node, cst.Call) and isinstance(
            decorator_node.func, cst.Name
        ):
            return decorator_node.func.value == self.decorator_name
        return False


class ProfileEnableTransformer(cst.CSTTransformer):
    """Inserts ``codeflash_line_profile.enable(...)`` after its import."""

    def __init__(self, filename: str) -> None:
        """Initialize with the filename for the enable() call."""
        self.found_import = False
        self.import_indentation: str | None = None
        self.filename = filename

    def leave_ImportFrom(  # noqa: N802
        self,
        original_node: cst.ImportFrom,
        updated_node: cst.ImportFrom,
    ) -> cst.ImportFrom:
        """Detect the line_profiler import statement."""
        if isinstance(original_node.names, cst.ImportStar):
            return updated_node
        if (
            isinstance(original_node.module, cst.Name)
            and original_node.module.value == "line_profiler"
            and any(
                name.name.value == "profile"
                and (
                    not name.asname
                    or (
                        isinstance(name.asname, cst.AsName)
                        and isinstance(name.asname.name, cst.Name)
                        and name.asname.name.value == "codeflash_line_profile"
                    )
                )
                for name in original_node.names
            )
        ):
            self.found_import = True
            if hasattr(original_node, "leading_lines"):
                leading_whitespace = (
                    original_node.leading_lines[-1].whitespace
                    if original_node.leading_lines
                    else ""
                )
                self.import_indentation = leading_whitespace

        return updated_node

    def leave_Module(  # noqa: N802
        self, original_node: cst.Module, updated_node: cst.Module
    ) -> cst.Module:
        """Insert the enable() call after the line_profiler import."""
        if not self.found_import:
            return updated_node

        new_body = list(updated_node.body)

        import_index = None
        for i, stmt in enumerate(new_body):
            if isinstance(stmt, cst.SimpleStatementLine):
                for small_stmt in stmt.body:
                    if (
                        isinstance(small_stmt, cst.ImportFrom)
                        and not isinstance(small_stmt.names, cst.ImportStar)
                        and isinstance(small_stmt.module, cst.Name)
                        and small_stmt.module.value == "line_profiler"
                        and any(
                            name.name.value == "profile"
                            and (
                                not name.asname
                                or (
                                    isinstance(name.asname, cst.AsName)
                                    and isinstance(name.asname.name, cst.Name)
                                    and name.asname.name.value
                                    == "codeflash_line_profile"
                                )
                            )
                            for name in small_stmt.names
                        )
                    ):
                        import_index = i
                        break
                if import_index is not None:
                    break

        if import_index is not None:
            enable_statement = cst.parse_statement(
                f"codeflash_line_profile.enable(output_prefix='{self.filename}')"
            )
            new_body.insert(import_index + 1, enable_statement)

        return updated_node.with_changes(body=new_body)


class LineProfilerImportAdder(cst.CSTTransformer):
    """Adds an import statement to a module if not already present."""

    def __init__(self, import_statement: str) -> None:
        """Initialize with the import statement to add."""
        self.import_statement = import_statement
        self.has_import = False

    def leave_Module(  # noqa: N802
        self, original_node: cst.Module, updated_node: cst.Module
    ) -> cst.Module:
        """Prepend the import statement if not already present."""
        if self.has_import:
            return updated_node
        import_node = cst.parse_statement(self.import_statement)
        return updated_node.with_changes(
            body=[import_node, *list(updated_node.body)]
        )

    def visit_ImportFrom(self, node: cst.ImportFrom) -> None:  # noqa: N802
        """Check if line_profiler profile is already imported."""
        if isinstance(node.names, cst.ImportStar):
            return
        if (
            node.module
            and isinstance(node.module, cst.Name)
            and node.module.value == "line_profiler"
        ):
            for import_alias in node.names:
                if import_alias.name.value == "profile":
                    self.has_import = True


def add_decorator_to_qualified_function(
    module: cst.Module, qualified_name: str, decorator_name: str
) -> cst.Module:
    """Apply *decorator_name* to the function matching *qualified_name*."""
    transformer = LineProfilerDecoratorAdder(qualified_name, decorator_name)
    return module.visit(transformer)


def add_profile_enable(
    original_code: str, line_profile_output_file: str
) -> str:
    """Insert ``codeflash_line_profile.enable(...)`` after its import."""
    module = cst.parse_module(original_code)
    transformer = ProfileEnableTransformer(line_profile_output_file)
    modified_module = module.visit(transformer)
    return modified_module.code


def add_decorator_imports(
    function_to_optimize: FunctionToOptimize,
    helper_functions: Sequence[FunctionSource],
) -> Path:
    """Add line_profiler decorators and imports to the target and helper files.

    Returns the path to the baseline line-profile output file.
    """
    file_paths: dict[Path, list[str]] = defaultdict(list)
    line_profile_output_file = get_run_tmp_file(Path("baseline_lprof"))
    file_paths[function_to_optimize.file_path].append(
        function_to_optimize.qualified_name
    )
    for elem in helper_functions:
        file_paths[elem.file_path].append(elem.qualified_name)
    for file_path, fns_present in file_paths.items():
        file_contents = file_path.read_text("utf-8")
        module_node = cst.parse_module(file_contents)
        for fn_name in fns_present:
            module_node = add_decorator_to_qualified_function(
                module_node, fn_name, "codeflash_line_profile"
            )
        transformer = LineProfilerImportAdder(
            "from line_profiler import profile as codeflash_line_profile"
        )
        module_node = module_node.visit(transformer)
        modified_code = module_node.code
        with file_path.open("w", encoding="utf-8") as file:
            file.write(modified_code)
    file_contents = function_to_optimize.file_path.read_text("utf-8")
    modified_code = add_profile_enable(
        file_contents, line_profile_output_file.as_posix()
    )
    function_to_optimize.file_path.write_text(modified_code, "utf-8")
    return line_profile_output_file
