"""Jedi-based function reference resolution for optimization review."""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING, Any

import attrs

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize

log = logging.getLogger(__name__)

_MAX_CONTEXT_LEN = 8000


@attrs.frozen
class ReferenceInfo:
    """A single call-site reference to a function."""

    file_path: Path
    line: int
    column: int
    context: str
    reference_type: str
    caller_function: str | None = None


def _find_function_position(
    names: list[Any],
    function_name: str,
    class_name: str | None,
) -> tuple[int, int] | None:
    """Locate the definition position of *function_name* in Jedi names."""
    for name in names:
        if name.type != "function" or name.name != function_name:
            continue
        if class_name:
            parent = name.parent()
            if parent and parent.name == class_name and parent.type == "class":
                return (name.line, name.column)
        else:
            return (name.line, name.column)
    return None


def _process_reference(
    ref: Any,
    function_file: Path,
    function_pos: tuple[int, int],
    tests_root: Path | None,
    seen: set[tuple[Any, int, int]],
) -> ReferenceInfo | None:
    """Process a single Jedi reference into a ReferenceInfo, or None."""
    if not ref.module_path:
        return None

    from pathlib import Path as _Path  # noqa: PLC0415

    ref_path = _Path(ref.module_path)

    # Skip the definition itself.
    if ref_path == function_file and ref.line == function_pos[0]:
        return None

    # Skip test files.
    if tests_root:
        try:
            ref_path.relative_to(tests_root)
        except ValueError:
            pass
        else:
            return None

    loc = (ref_path, ref.line, ref.column)
    if loc in seen:
        return None
    seen.add(loc)

    # Context line.
    try:
        lines = ref_path.read_text(encoding="utf-8").splitlines()
        ctx = lines[ref.line - 1] if ref.line <= len(lines) else ""
    except Exception:  # noqa: BLE001
        ctx = ""

    # Determine caller function.
    caller: str | None = None
    try:
        parent = ref.parent()
        if parent and parent.type == "function":
            caller = parent.name
    except Exception:  # noqa: BLE001
        log.debug("Error determining caller for ref at %s", loc)

    return ReferenceInfo(
        file_path=ref_path,
        line=ref.line,
        column=ref.column,
        context=ctx.strip(),
        reference_type="call",
        caller_function=caller,
    )


def find_function_references(
    function: FunctionToOptimize,
    project_root: Path,
    tests_root: Path | None = None,
) -> list[ReferenceInfo]:
    """Find call sites for *function* across the project using Jedi.

    Excludes references inside the tests directory and the
    function's own definition.
    """
    import jedi  # type: ignore[import-untyped]  # noqa: PLC0415

    try:
        source = function.file_path.read_text(encoding="utf-8")
        script = jedi.Script(code=source, path=function.file_path)
        names = script.get_names(all_scopes=True, definitions=True)

        function_pos = _find_function_position(
            names,
            function.function_name,
            function.class_name,
        )
        if function_pos is None:
            return []

        project = jedi.Project(path=project_root)
        script = jedi.Script(
            code=source,
            path=function.file_path,
            project=project,
        )
        references = script.get_references(
            line=function_pos[0],
            column=function_pos[1],
        )

        result: list[ReferenceInfo] = []
        seen: set[tuple[Any, int, int]] = set()
        for ref in references:
            info = _process_reference(
                ref,
                function.file_path,
                function_pos,
                tests_root,
                seen,
            )
            if info is not None:
                result.append(info)
    except Exception:  # noqa: BLE001
        log.debug(
            "Error finding references for %s",
            function.function_name,
            exc_info=True,
        )
        return []
    else:
        return result


def _extract_calling_function_source(
    source_code: str,
    function_name: str,
    ref_line: int,
) -> str | None:
    """Extract the source of the function that contains *ref_line*."""
    try:
        lines = source_code.splitlines()
        tree = ast.parse(source_code)
        for node in ast.walk(tree):
            if (
                isinstance(
                    node,
                    (ast.FunctionDef, ast.AsyncFunctionDef),
                )
                and node.name == function_name
            ):
                end = node.end_lineno or node.lineno
                if node.lineno <= ref_line <= end:
                    return "\n".join(lines[node.lineno - 1 : end])
    except Exception:  # noqa: BLE001
        return None
    return None


def _collect_file_contexts(
    file_refs: list[ReferenceInfo],
    content: str,
) -> list[str]:
    """Collect calling-function context snippets for refs in one file."""
    lines = content.splitlines()
    callers_seen: set[str] = set()
    contexts: list[str] = []
    for ref in file_refs:
        caller = ref.caller_function or "<module>"
        if caller in callers_seen:
            continue
        callers_seen.add(caller)

        if ref.caller_function:
            code = _extract_calling_function_source(
                content,
                ref.caller_function,
                ref.line,
            )
            if code:
                contexts.append(code)
        else:
            start = max(0, ref.line - 3)
            end = min(len(lines), ref.line + 2)
            snippet = "\n".join(lines[start:end])
            contexts.append(snippet)
    return contexts


def format_references_as_markdown(
    references: list[ReferenceInfo],
    source_file: Path,
    project_root: Path,
) -> str:
    """Format *references* as markdown code blocks with calling function context."""
    refs_by_file: dict[Path, list[ReferenceInfo]] = {}
    for ref in references:
        if ref.file_path == source_file and ref.reference_type in (
            "import",
            "reexport",
        ):
            continue
        refs_by_file.setdefault(ref.file_path, []).append(ref)

    output = ""
    context_len = 0
    for ref_file, file_refs in refs_by_file.items():
        if context_len > _MAX_CONTEXT_LEN:
            break
        try:
            rel = ref_file.relative_to(project_root)
        except ValueError:
            continue
        try:
            content = ref_file.read_text(encoding="utf-8")
        except Exception:  # noqa: BLE001
            log.debug("Cannot read %s for reference context", ref_file)
            continue

        contexts = _collect_file_contexts(file_refs, content)
        if contexts:
            output += f"```python:{rel.as_posix()}\n"
            output += "\n".join(contexts)
            output += "\n```\n"
            context_len += sum(len(c) for c in contexts)

    return output
