from __future__ import annotations

import ast
import base64
import importlib
import inspect
import os
import sys
import textwrap
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Optional, List, Mapping

import dill


# ---------- small data structs -------------------------------------------------


@dataclass(frozen=True)
class DependencyInfo:
    """
    Single dependency entry.

    Example:
        DependencyInfo(
            root_module="databricks",
            submodule="databricks.sdk",
            root_path="/path/to/databricks"
        )
    """

    root_module: str
    submodule: str
    root_path: Optional[str]


@dataclass(frozen=True)
class DependencyCheckResult:
    """
    Result of checking whether a dependency is importable in the current env.
    """

    root_module: str
    root_path: Optional[str]
    importable: bool
    runtime_file: Optional[str]
    error: Optional[str]


# ---------- helpers ------------------------------------------------------------


def _find_package_root_from_file(module_file: str) -> str:
    """
    Given a module __file__, find the top-most package directory.

    Walk upwards from the directory containing the file as long
    as there is an __init__.py present. Return that top-most dir.
    """
    path = os.path.abspath(module_file)
    current_dir = os.path.dirname(path)

    last_pkg_dir: Optional[str] = None

    # If the current dir itself is a package, mark it
    if os.path.isfile(os.path.join(current_dir, "__init__.py")):
        last_pkg_dir = current_dir
    else:
        return module_file

    # Walk upwards while parent is also a package
    while True:
        parent = os.path.dirname(current_dir)
        if parent == current_dir:  # FS root
            break

        init_path = os.path.join(parent, "__init__.py")
        if os.path.isfile(init_path):
            last_pkg_dir = parent
            current_dir = parent
        else:
            break

    return last_pkg_dir or os.path.dirname(path)


def _extract_function_source(
    raw_src: str,
    qualname: str,
    func_name: str,
) -> str:
    """
    Given raw source returned by inspect.getsource(func), try to extract just the
    function definition for the target function.

    Supports:
      - top-level functions: qualname == name
      - nested/local functions: qualname like "outer.<locals>.inner"

    For nested funcs, we *try* to walk the AST and use ast.get_source_segment()
    to pull the exact code for the target node. If that fails (e.g. because
    raw_src is just an indented inner function and ast.parse blows up with
    IndentationError), we fall back to simply dedenting the whole block.
    """
    # strip out "<locals>" segments from qualname
    parts = [p for p in qualname.split(".") if p != "<locals>"]

    # Simple case: top-level function
    if len(parts) == 1 and parts[0] == func_name:
        return textwrap.dedent(raw_src)

    # Nested/local function: try AST-based extraction
    try:
        tree = ast.parse(raw_src)
    except (SyntaxError, IndentationError):
        # Most likely case: raw_src is just an indented inner def, like:
        # "        def inner(x):\n            return x + 1\n"
        # Dedent and treat that as the function source.
        return textwrap.dedent(raw_src)

    target_path = parts  # e.g. ["outer", "inner"]

    def _walk(node: ast.AST, idx: int) -> Optional[ast.AST]:
        body = getattr(node, "body", [])
        for child in body:
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                if child.name == target_path[idx]:
                    if idx == len(target_path) - 1:
                        return child
                    return _walk(child, idx + 1)
        return None

    found_node = _walk(tree, 0)

    if not found_node:
        # fallback: give the whole block; better than blowing up
        return textwrap.dedent(raw_src)

    segment = ast.get_source_segment(raw_src, found_node)
    if segment is None:
        return textwrap.dedent(raw_src)

    return textwrap.dedent(segment)


def _infer_raw_dependencies_and_aliases(
    func: Callable[..., Any]
) -> tuple[Dict[str, Optional[str]], Dict[str, str]]:
    """
    Infer:

      - raw_deps: { full_module_name -> original_module_file_or_None }
      - aliases:  { global_name_used_in_func -> full_module_name }

    by inspecting closure globals/nonlocals for module objects.
    Example alias:
        MATH_MOD -> "math"
        JSON_MOD -> "json"
    """
    raw_deps: Dict[str, Optional[str]] = {}
    aliases: Dict[str, str] = {}

    closure = inspect.getclosurevars(func)

    def _collect(mapping: Mapping[str, Any]) -> None:
        for name, val in mapping.items():
            if inspect.ismodule(val):
                mod_name = getattr(val, "__name__", None)
                if not mod_name:
                    continue
                mod_file = getattr(val, "__file__", None)
                if mod_name not in raw_deps or (not raw_deps[mod_name] and mod_file):
                    raw_deps[mod_name] = mod_file
                # record alias (e.g. MATH_MOD -> "math")
                aliases.setdefault(name, mod_name)

    _collect(closure.globals)
    _collect(closure.nonlocals)

    return raw_deps, aliases


def _build_dependency_infos(
    raw_deps: Mapping[str, Optional[str]],
) -> List[DependencyInfo]:
    """
    Turn raw deps ({full_module_name -> file}) into DependencyInfo objects:

        "databricks.sdk.service.compute" ->
            DependencyInfo(
                root_module="databricks",
                submodule="databricks.sdk",
                root_path="/path/to/databricks"
            )
    """
    infos: List[DependencyInfo] = []
    seen: set[tuple[str, str]] = set()

    for full_name in raw_deps.keys():
        parts = full_name.split(".")
        root = parts[0]

        if len(parts) > 1:
            sub = ".".join(parts[:2])
        else:
            sub = root

        key = (root, sub)
        if key in seen:
            continue
        seen.add(key)

        root_path: Optional[str] = None
        try:
            root_mod = importlib.import_module(root)
            root_file = getattr(root_mod, "__file__", None)
            if root_file:
                root_path = os.path.dirname(os.path.abspath(root_file))
        except Exception:
            root_path = None

        infos.append(DependencyInfo(root_module=root, submodule=sub, root_path=root_path))

    return infos


# ---------- main class --------------------------------------------------------


@dataclass
class SerializedFunction:
    """
    Function + its source code embedded as text.

    - Can be serialized (pickle/dill/msgpack/etc.)
    - Can be rebuilt on another Python process via exec()
    - Is directly callable: embedded_fn(*args, **kwargs)
    - Tracks dependency infos:
        DependencyInfo(root_module, submodule, root_path)
    - Captures the package root of the defining module (if resolvable)
    - Supports top-level and local/nested functions, as long as they do NOT
      capture nonlocal variables in their closure.
    """

    name: str
    source: str
    source_file_content: Optional[str]

    # List of DependencyInfo entries
    dependencies_map: List[DependencyInfo] = field(default_factory=list)

    # Global-name -> full module name alias mapping (e.g. "MATH_MOD" -> "math")
    aliases: Dict[str, str] = field(default_factory=dict)

    # Filesystem path to the top-most package containing the defining module
    package_root: Optional[str] = None

    # internal cache for the rebuilt function
    _compiled: Callable[..., Any] | None = field(default=None, init=False, repr=False)

    # ---------- constructors ----------
    @classmethod
    def _infer_imported_modules_from_file(
        cls,
        module_file: str,
        src: Optional[str] = None,
        package_root: Optional[str] = None,
    ) -> List[DependencyInfo]:
        """
        Given a Python source file, parse its imports and return DependencyInfo
        objects for modules that look relevant.

        Returns:
            List[DependencyInfo]
        """
        if not src:
            with open(module_file, "r", encoding="utf-8") as f:
                src = f.read()

        tree = ast.parse(src, filename=module_file)

        imported_module_names: set[str] = set()

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                # import x, import x.y.z
                for alias in node.names:
                    imported_module_names.add(alias.name)
            elif isinstance(node, ast.ImportFrom):
                # from x import y; from x.y import z; from .x import y
                if node.module is None:
                    # "from . import foo" -> skip unless you want to get fancy
                    continue
                full = ("." * node.level + node.module) if node.level else node.module
                imported_module_names.add(full)

        # Resolve names → real modules, filter by package_root if given
        resolved_names: set[str] = set()

        for name in imported_module_names:
            try:
                # Handle relative names relative to this module's package
                if name.startswith("."):
                    if not package_root:
                        continue
                    rel_path = os.path.relpath(module_file, package_root)
                    # e.g. my_pkg/mod.py -> my_pkg.mod
                    pkg = os.path.splitext(rel_path)[0].replace(os.path.sep, ".")
                    pkg = pkg.rsplit(".", 1)[0]  # strip final segment
                    full_name = pkg + name  # let importlib normalize it
                    spec = importlib.util.find_spec(full_name)
                else:
                    spec = importlib.util.find_spec(name)

                if spec is None or spec.origin is None:
                    continue

                resolved_names.add(spec.name)
            except Exception:
                # noisy/edge imports shouldn't break everything
                continue

        # Now build DependencyInfo list reusing your existing helper
        raw_deps = {full_name: None for full_name in resolved_names}
        return _build_dependency_infos(raw_deps)

    @classmethod
    def from_callable(cls, func: Callable[..., Any]) -> "SerializedFunction":
        """
        Build an EmbeddedFunction from a function.

        Supports:
          - top-level functions
          - local/nested functions (qualname may contain '<locals>')

        Rejects:
          - bound methods (inspect.ismethod)
          - functions that capture nonlocal variables in their closure
        """
        # 1) Reject bound methods
        if inspect.ismethod(func):
            raise ValueError("EmbeddedFunction does not support bound methods")

        # 2) Qualname is how we identify nested/local defs
        qualname = getattr(func, "__qualname__", None)
        if not qualname:
            raise ValueError(f"Cannot derive qualname for {func!r}")

        # 4) Extract raw source for the function
        try:
            raw_src = inspect.getsource(func)
        except OSError as exc:
            raise ValueError(f"Cannot extract source for {func!r}: {exc}") from exc

        # 5) Trim down to just the def block for this specific function
        src = _extract_function_source(raw_src, qualname, func.__name__)

        # 6) Infer raw deps (module_name -> file_or_None) + aliases
        raw_deps, aliases = _infer_raw_dependencies_and_aliases(func)

        # 7) Add the defining module; track package_root
        module_name = getattr(func, "__module__", None)
        package_root: Optional[str] = None
        module_file: Optional[str] = None
        source_file_content: Optional[str] = None

        if module_name:
            mod = sys.modules.get(module_name)
            if mod is not None:
                module_file = getattr(mod, "__file__", None)
                if module_file:
                    module_file = os.path.abspath(module_file)
                    package_root = _find_package_root_from_file(module_file)
                    # Always include the defining module itself
                    raw_deps.setdefault(module_name, module_file)

        # 8) Convert closure/module deps -> DependencyInfo list
        base_dep_infos: List[DependencyInfo] = _build_dependency_infos(raw_deps)

        # 9) Add imported modules from the defining file (also DependencyInfo)
        imported_dep_infos: List[DependencyInfo] = []
        if module_file:
            with open(module_file, "r", encoding="utf-8") as f:
                source_file_content = f.read()
            imported_dep_infos = cls._infer_imported_modules_from_file(
                module_file=module_file,
                src=source_file_content,
                package_root=package_root,
            )

        # 10) Merge and dedupe DependencyInfo entries
        merged: dict[tuple[str, str], DependencyInfo] = {}

        for info in base_dep_infos + imported_dep_infos:
            key = (info.root_module, info.submodule)
            # last-one-wins; they're usually identical anyway
            merged[key] = info

        dep_infos = list(merged.values())

        # 11) Build the EmbeddedFunction dataclass
        built = cls(
            name=func.__name__,
            source=src,
            source_file_content=source_file_content,
            dependencies_map=dep_infos,
            aliases=aliases,
            package_root=package_root,
        )

        # Cache the original callable for local use
        built._compiled = func

        return built

    # ---------- rebuild ----------
    def build(self, globals_ns: Dict[str, Any] | None = None) -> Callable[..., Any]:
        """
        Rebuild the function object from stored source.

        globals_ns:
            Optional globals dict to exec into. If None, a fresh empty dict
            is used (so you must import everything you need inside the function).

        We also:
          - add dependency root paths to sys.path (if they exist),
          - restore module aliases (e.g. MATH_MOD -> math) based on the
            original closure's module-valued globals.
        """
        if globals_ns is None:
            globals_ns = {}

        # 1) Make sure dependency roots are on sys.path
        import sys as _sys
        import os as _os

        for dep in self.dependencies_map:
            root_path = dep.root_path
            if (
                root_path
                and _os.path.isdir(root_path)
                and root_path not in _sys.path
            ):
                # prepend so user code takes precedence over site-packages
                _sys.path.insert(0, root_path)

        # 2) Restore module aliases first so the function body can use them
        for alias, module_name in self.aliases.items():
            try:
                globals_ns[alias] = importlib.import_module(module_name)
            except Exception:
                # If import fails, leave the alias missing; calling the function
                # will raise NameError, which is explicit enough.
                pass

        # 3) Exec the function source
        exec(self.source, globals_ns)  # noqa: S102 - intentional

        fn = globals_ns.get(self.name)
        if not callable(fn):
            raise RuntimeError(
                f"After exec, object named {self.name!r} is not callable "
                f"(got {type(fn)!r})"
            )
        return fn

    def _ensure_compiled(self) -> Callable[..., Any]:
        """
        Lazily compile the function once and cache it.
        """
        if self._compiled is None:
            self._compiled = self.build()
        return self._compiled

    @property
    def func(self):
        return self._ensure_compiled()

    # ---------- make it callable ----------

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """
        Call the embedded function as if it were a normal Python function.
        """
        fn = self._ensure_compiled()
        return fn(*args, **kwargs)

    # ---------- dependency inspection ----------

    def list_dependencies(self) -> List[DependencyInfo]:
        """
        Return the dependency infos:
            DependencyInfo(root_module, submodule, root_path)
        """
        return list(self.dependencies_map)

    def check_dependencies(self) -> Dict[str, DependencyCheckResult]:
        """
        Check whether dependencies are importable *on this machine* and where
        they resolve to.

        Uses the submodule (e.g. "databricks.sdk") as the thing to import.

        Returns:
            {
              submodule: DependencyCheckResult(...),
              ...
            }
        """
        results: Dict[str, DependencyCheckResult] = {}

        for info in self.dependencies_map:
            try:
                mod = importlib.import_module(info.submodule)
                importable = True
                runtime_file = getattr(mod, "__file__", None)
                error: Optional[str] = None
            except Exception as exc:
                importable = False
                runtime_file = None
                error = repr(exc)

            results[info.submodule] = DependencyCheckResult(
                root_module=info.root_module,
                root_path=info.root_path,
                importable=importable,
                runtime_file=runtime_file,
                error=error,
            )

        return results

    def to_command(
        self,
        result_tag: str,
        args: list = None,
        kwargs: dict = None,
        env_keys: Optional[List[str]] = None,
        use_dill: bool = False,

    ) -> str:
        """
        Build a Python command string suitable for Databricks `command_execution`.

        The generated code will:
          - best-effort add dependency root paths to sys.path
          - restore module aliases used by the function
          - reconstruct args/kwargs via dill (sent from driver)
          - call the function with *args / **kwargs
          - capture stdout/stderr produced during the call
          - print a base64-encoded dill-serialized dict:
                {
                  "result_b64": <base64 of dill(result or DataFrame(result))>,
                  "stdout": "<captured stdout>",
                  "stderr": "<captured stderr>",
                }

        Args:
            args: positional args to call the function with.
            kwargs: keyword args to call the function with.
            env_keys: optional environment variables to forward.
            use_dill: when True, serialize the function itself with dill rather than
                embedding its source code. Defaults to False to keep the existing
                source-embedding behaviour.
            result_tag: optional marker string to wrap the printed payload. When
                provided, the payload is printed as ``<result_tag><json><result_tag>``
                to simplify parsing downstream.
        """
        import dill  # driver-side
        import base64 as _b64
        import os as _os

        if self.package_root:
            roots = [_os.path.dirname(self.package_root)]
        else:
            roots = []

        def _b64_dumps(obj: Any) -> str:
            return _b64.b64encode(dill.dumps(obj)).decode("utf-8")

        func_ser: Optional[str] = None
        if use_dill:
            func_ser = _b64_dumps(self._ensure_compiled())

        imports = set(d.submodule for d in self.dependencies_map)

        lines: list[str] = [
            "import sys",
            "import os",
            "import base64",
            "import importlib",
            "import json",
            "import dill",
            "import pandas",
            "import io",
            "import contextlib",
        ]

        # --- embed env vars into remote process ---
        if env_keys:
            env_vars = {
                k: os.getenv(k)
                for k in env_keys
                if os.getenv(k) is not None
            }

            if env_vars:
                # env_vars is a plain dict on the driver; we stringify it into code
                lines.append(f"__embedded_env = {env_vars!r}")
                lines.extend([
                    "for __k, __v in __embedded_env.items():",
                    "    if __v is not None:",
                    "        os.environ[__k] = __v",
                ])

        # --- sys.path bootstrapping on remote ---
        if roots:
            lines.append(f"for __p in {roots!r}:")
            lines.extend([
                "    if __p and __p not in sys.path:",
                "        if os.path.isdir(__p):",
                "            sys.path.insert(0, __p)",
            ])

        # --- best-effort import of dependencies on remote ---
        if imports:
            lines.extend([
                "try:\n    import %s\nexcept:\n    pass" % _ for _ in imports
            ])

        # --- restore module aliases on remote ---
        if self.aliases:
            lines.append(f"__embedded_aliases = {self.aliases!r}")
            lines.append("for __name, __mod in __embedded_aliases.items():")
            lines.append("    globals()[__name] = importlib.import_module(__mod)")

        # Serialize args/kwargs as plain data
        args_ser = _b64_dumps(tuple(args or []))
        kwargs_ser = _b64_dumps(dict(kwargs or {}))

        if use_dill:
            lines.append(
                f"_embedded_func = dill.loads(base64.b64decode({func_ser!r}.encode('utf-8')))"
            )
        else:
            # Prefer full file content if we have it, otherwise just the function source
            source_block = self.source_file_content or self.source
            lines.extend([
                "",
                source_block,
                "",
                f"_embedded_func = {self.name}",
            ])

        # --- remote-side deserialization + log capture ---
        lines.extend([
            # deserialize args / kwargs
            f"_embedded_args = dill.loads(base64.b64decode({args_ser!r}.encode('utf-8')))",
            f"_embedded_kwargs = dill.loads(base64.b64decode({kwargs_ser!r}.encode('utf-8')))",
            "_embedded_result = _embedded_func(*_embedded_args, **_embedded_kwargs)",
            "_ser_result = dill.dumps(_embedded_result)",
        ])

        lines.append(
            f"print({result_tag!r} + base64.b64encode(_ser_result).decode('utf-8') + {result_tag!r})"
        )

        return "\n".join(lines)

    def parse_command_result(self, result: str, raise_error: bool = True):
        """
        Decode + deserialize remote result.

        On success:
            returns whatever the remote function produced.

        On remote error:
            expects a dict payload:
                {
                    'type': <exception class name>,
                    'message': <str(e)>,
                    'traceback': <formatted traceback string>,
                }
            and raises RemoteExecutionError.
        """
        raw = base64.b64decode(result)
        obj = dill.loads(raw)

        return obj
