"""Extract source code for discovered functions."""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize

log = logging.getLogger(__name__)


def extract_function_source(
    function: FunctionToOptimize,
) -> str:
    """Read the source text of *function* from disk."""
    lines = function.file_path.read_text().splitlines(keepends=True)
    start = function.starting_line
    end = function.ending_line
    if start is None or end is None:
        msg = (
            f"Cannot extract source for"
            f" {function.qualified_name}:"
            f" missing line numbers"
        )
        raise ValueError(msg)
    return "".join(lines[start - 1 : end])


def get_code(
    functions_to_optimize: list[FunctionToOptimize],
) -> tuple[str | None, set[tuple[str, str]]]:
    """Return source code for function(s), plus dunder methods in their class.

    *functions_to_optimize* is either a single module-level function /
    class method, or multiple methods of the **same** class.

    Returns ``(source_code, contextual_dunder_methods)`` where
    *contextual_dunder_methods* is a set of ``(class_name, dunder_name)``
    tuples for dunder methods found in the same class.
    """
    if (
        not functions_to_optimize
        or (
            functions_to_optimize[0].parents
            and functions_to_optimize[0].parents[0].type != "ClassDef"
        )
        or (
            len(functions_to_optimize[0].parents) > 1
            or (
                (len(functions_to_optimize) > 1)
                and len(
                    {fn.parents[0] for fn in functions_to_optimize},
                )
                != 1
            )
        )
    ):
        return None, set()

    file_path: Path = functions_to_optimize[0].file_path
    class_skeleton: set[tuple[int, int | None]] = set()
    contextual_dunder_methods: set[tuple[str, str]] = set()
    target_code: str = ""

    def find_target(
        node_list: list[ast.stmt],
        name_parts: tuple[str, ...],
    ) -> ast.AST | None:
        target: (
            ast.FunctionDef
            | ast.AsyncFunctionDef
            | ast.ClassDef
            | ast.Assign
            | ast.AnnAssign
            | None
        ) = None
        for node in node_list:
            if (
                isinstance(
                    node,
                    (
                        ast.FunctionDef,
                        ast.AsyncFunctionDef,
                        ast.ClassDef,
                    ),
                )
                and node.name == name_parts[0]
            ):
                target = node
                break
            if (
                isinstance(node, ast.Assign)
                and len(node.targets) == 1
                and isinstance(node.targets[0], ast.Name)
                and node.targets[0].id == name_parts[0]
            ) or (
                isinstance(node, ast.AnnAssign)
                and hasattr(node.target, "id")
                and node.target.id == name_parts[0]
            ):
                if class_skeleton:
                    break
                target = node
                break

        if target is None or len(name_parts) == 1:
            return target

        if not isinstance(target, ast.ClassDef) or len(name_parts) < 2:
            return None
        method_name: str = name_parts[1]
        class_skeleton.add(
            (target.lineno, target.body[0].lineno - 1),
        )
        cbody = target.body
        if isinstance(cbody[0], ast.expr):
            class_skeleton.add((cbody[0].lineno, cbody[0].end_lineno))
            cbody = cbody[1:]
        for cnode in cbody:
            cnode_name: str
            if (
                isinstance(
                    cnode,
                    (ast.FunctionDef, ast.AsyncFunctionDef),
                )
                and len(cnode_name := cnode.name) > 4
                and cnode_name != method_name
                and cnode_name.isascii()
                and cnode_name.startswith("__")
                and cnode_name.endswith("__")
            ):
                contextual_dunder_methods.add(
                    (target.name, cnode_name),
                )
                class_skeleton.add(
                    (cnode.lineno, cnode.end_lineno),
                )

        return find_target(target.body, (method_name,))

    with file_path.open(encoding="utf8") as file:
        source_code: str = file.read()
    try:
        module_node: ast.Module = ast.parse(source_code)
    except SyntaxError:
        log.exception("get_code - Syntax error while parsing code")
        return None, set()

    lines: list[str] = source_code.splitlines(keepends=True)
    if len(functions_to_optimize[0].parents) == 1:
        if functions_to_optimize[0].parents[0].type == "ClassDef":
            qualified_name_parts_list: list[tuple[str, str] | tuple[str]] = [
                (fto.parents[0].name, fto.function_name)
                for fto in functions_to_optimize
            ]
        else:
            log.error(
                "get_code does not support inner functions: %s",
                functions_to_optimize[0].parents,
            )
            return None, set()
    elif len(functions_to_optimize[0].parents) == 0:
        qualified_name_parts_list = [
            (functions_to_optimize[0].function_name,),
        ]
    else:
        log.error(
            "get_code does not support more than one level of"
            " nesting. Parents: %s",
            functions_to_optimize[0].parents,
        )
        return None, set()

    for qualified_name_parts in qualified_name_parts_list:
        target_node = find_target(
            module_node.body,
            qualified_name_parts,
        )
        if target_node is None:
            continue
        if not isinstance(
            target_node,
            (
                ast.FunctionDef,
                ast.AsyncFunctionDef,
                ast.ClassDef,
                ast.Assign,
                ast.AnnAssign,
            ),
        ):
            continue

        if (
            isinstance(
                target_node,
                (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef),
            )
            and target_node.decorator_list
        ):
            target_code += "".join(
                lines[
                    target_node.decorator_list[0].lineno
                    - 1 : target_node.end_lineno
                ],
            )
        else:
            target_code += "".join(
                lines[target_node.lineno - 1 : target_node.end_lineno],
            )

    if not target_code:
        return None, set()
    class_list: list[tuple[int, int | None]] = sorted(class_skeleton)
    class_code = "".join(
        [
            "".join(lines[s_lineno - 1 : e_lineno])
            for (s_lineno, e_lineno) in class_list
        ],
    )
    return class_code + target_code, contextual_dunder_methods
