"""Cache key derivation: transitive AST normalization + blake3.

Walks the call graph to depth `max_depth` (default 5), hashing the
normalized AST of the function and all reachable callees. Caches the
composite function identity hash, invalidating on source file mtime change.
"""

import ast
import inspect
import os
import pickle
import textwrap
from typing import Dict, Optional, Sequence, Tuple

try:
    from blake3 import blake3 as _blake3
    def _hash(data: bytes) -> bytes:
        return _blake3(data).digest()
except ImportError:
    import hashlib
    def _hash(data: bytes) -> bytes:
        return hashlib.blake2b(data, digest_size=32).digest()

# Cache: fn_id -> (mtime_tuple, identity_hash)
_identity_cache: Dict[int, Tuple[tuple, str]] = {}


def make_cache_key(
    fn,
    args: tuple,
    kwargs: dict,
    ignore: Optional[Sequence[str]] = None,
    tag: str = "",
    max_depth: int = 5,
) -> str:
    """Produce a hex 32-byte cache key for fn(*args, **kwargs)."""
    fn_identity = _cached_identity(fn, max_depth)
    filtered = _filter_args(fn, args, kwargs, ignore)
    blob = (
        fn_identity.encode("utf-8")
        + b"\x00"
        + pickle.dumps(filtered, protocol=5)
    )
    key_hex = _hash(blob).hex()
    if tag:
        return f"{tag}/{key_hex}"
    return key_hex


def _cached_identity(fn, max_depth: int) -> str:
    """Return mtime-cached composite AST identity for fn + callees."""
    fn_id = id(fn)
    mtimes = _source_mtimes(fn, max_depth)

    cached = _identity_cache.get(fn_id)
    if cached and cached[0] == mtimes:
        return cached[1]

    identity = _transitive_ast(fn, max_depth)
    _identity_cache[fn_id] = (mtimes, identity)
    return identity


def _source_mtimes(fn, max_depth: int) -> tuple:
    """Collect mtime of source files for fn and reachable callees."""
    seen = set()
    mtimes = []
    _collect_mtimes(fn, max_depth, seen, mtimes)
    return tuple(sorted(mtimes))


def _collect_mtimes(fn, depth, seen, mtimes):
    if depth <= 0 or id(fn) in seen:
        return
    seen.add(id(fn))
    try:
        source_file = inspect.getfile(fn)
        if os.path.isfile(source_file):
            mtimes.append((source_file, os.path.getmtime(source_file)))
    except (TypeError, OSError):
        pass
    if depth > 1:
        for callee in _find_callees(fn):
            _collect_mtimes(callee, depth - 1, seen, mtimes)


def _transitive_ast(fn, max_depth: int) -> str:
    """Build composite AST identity: fn + all reachable callees."""
    parts = []
    seen = set()
    _collect_ast(fn, max_depth, seen, parts)
    return "\x00".join(parts)


def _collect_ast(fn, depth, seen, parts):
    if depth <= 0 or id(fn) in seen:
        return
    seen.add(id(fn))
    parts.append(_normalize_ast(fn))
    if depth > 1:
        for callee in _find_callees(fn):
            _collect_ast(callee, depth - 1, seen, parts)


def _normalize_ast(fn) -> str:
    """Return canonical, comment-free AST dump of fn."""
    try:
        source = textwrap.dedent(inspect.getsource(fn))
    except (OSError, TypeError):
        return getattr(fn, "__qualname__", getattr(fn, "__name__", repr(fn)))

    try:
        tree = ast.parse(source)
    except SyntaxError:
        return getattr(fn, "__qualname__", repr(fn))

    func = _find_funcdef(tree)
    if func is None:
        return ast.dump(tree)

    func.decorator_list = []
    _strip_locations(func)
    return ast.dump(func)


def _find_funcdef(tree):
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            return node
    return None


def _strip_locations(node):
    for child in ast.walk(node):
        for attr in ("lineno", "col_offset", "end_lineno", "end_col_offset", "type_comment"):
            try:
                delattr(child, attr)
            except AttributeError:
                pass


def _find_callees(fn):
    """Find callable objects referenced in fn's body via static analysis."""
    callees = []
    try:
        source = textwrap.dedent(inspect.getsource(fn))
        tree = ast.parse(source)
    except (OSError, TypeError, SyntaxError):
        return callees

    module = inspect.getmodule(fn)
    if module is None:
        return callees

    # Collect all Name nodes used in Call context
    for node in ast.walk(tree):
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
            name = node.func.id
            # Try to resolve in fn's globals, then module
            obj = None
            if hasattr(fn, "__globals__") and name in fn.__globals__:
                obj = fn.__globals__[name]
            elif hasattr(module, name):
                obj = getattr(module, name)
            if obj is not None and callable(obj) and inspect.isfunction(obj):
                callees.append(obj)

    return callees


def _filter_args(fn, args, kwargs, ignore):
    if not ignore:
        return (args, kwargs)
    try:
        sig = inspect.signature(fn)
        params = list(sig.parameters.keys())
    except (ValueError, TypeError):
        return (args, kwargs)
    filtered_args = tuple(
        val for i, val in enumerate(args)
        if i >= len(params) or params[i] not in ignore
    )
    filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ignore}
    return (filtered_args, filtered_kwargs)
