# yggdrasil/ser.py
from __future__ import annotations

import ast
import base64
import builtins
import inspect
import json
import os
import sys
import textwrap
import zlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, Optional, Tuple

import dill

PyVer = Tuple[int, int]


def _pyver() -> PyVer:
    v = sys.version_info
    return v.major, v.minor


def _b64e(b: bytes) -> str:
    return base64.b64encode(b).decode("ascii")


def _b64d(s: str) -> bytes:
    return base64.b64decode(s.encode("ascii"))


def _is_python_function(obj: Any) -> bool:
    return inspect.isfunction(obj) or inspect.ismethod(obj) or isinstance(obj, type(lambda: 0))


def _safe_dill_dumps(obj: Any) -> Optional[bytes]:
    try:
        return dill.dumps(obj, recurse=True)
    except Exception:
        return None


def _safe_dill_loads(b: Optional[bytes]) -> Any:
    if b is None:
        raise ValueError("No dill bytes to load")
    return dill.loads(b)


def _infer_package_root(fn: Callable[..., Any]) -> Tuple[str, ...]:
    """
    Return the *topmost* package directory for the module defining `fn`.

    Walk upwards from the module file directory while an __init__.py exists.
    Example:
      /repo/my_pkg/sub/mod.py  -> returns /repo/my_pkg
    If no __init__.py is found, falls back to the module's parent directory.
    """
    try:
        mod = inspect.getmodule(fn)
        file = getattr(mod, "__file__", None)
        if not file:
            return tuple()

        p = Path(file).resolve()
        if p.suffix.lower() != ".py":
            return tuple()

        cur = p.parent  # start at directory containing the module file

        # If this directory isn't a package at all, just return it
        if not (cur / "__init__.py").exists():
            return (str(cur),)

        # Climb while parent is still a package (has __init__.py)
        # Stop at the highest directory that is still a package.
        top = cur
        while True:
            parent = top.parent
            if parent == top:
                break
            if (parent / "__init__.py").exists():
                top = parent
                continue
            break

        return (str(top),)
    except Exception:
        return tuple()


def _capture_exec_env(fn: Callable[..., Any]) -> Dict[str, Dict[str, str]]:
    """
    Capture a minimal environment for exec(source) reconstruction.

    - globals: names referenced by fn.__code__.co_names that exist in fn.__globals__ (excluding builtins)
    - freevars: closure cell contents for fn.__code__.co_freevars

    Values are stored as base64(dill.dumps(value)) (best-effort).
    """
    code = fn.__code__
    fn_globals = getattr(fn, "__globals__", {}) or {}

    referenced = set(code.co_names or ())
    builtins_set = set(dir(builtins))

    g_payload: Dict[str, str] = {}
    for name in referenced:
        if name in builtins_set:
            continue
        if name in fn_globals:
            dumped = _safe_dill_dumps(fn_globals[name])
            if dumped is not None:
                g_payload[name] = _b64e(dumped)

    fv_payload: Dict[str, str] = {}
    freevars = code.co_freevars or ()
    closure = fn.__closure__ or ()
    if freevars and closure and len(freevars) == len(closure):
        for name, cell in zip(freevars, closure):
            try:
                val = cell.cell_contents
            except ValueError:
                continue
            dumped = _safe_dill_dumps(val)
            if dumped is not None:
                fv_payload[name] = _b64e(dumped)

    return {"globals": g_payload, "freevars": fv_payload}


def _capture_module_imports(fn: Callable[..., Any]) -> str:
    """
    Capture top-level imports from the module file where `fn` is defined.

    We extract only:
      - `import x`
      - `import x as y`
      - `from a.b import c`
      - `from a.b import c as d`

    Returned as a single string block (may be empty).
    """
    try:
        src_file = inspect.getsourcefile(fn) or inspect.getfile(fn)
        if not src_file:
            return ""
        p = Path(src_file)
        if not p.exists() or p.suffix.lower() != ".py":
            return ""

        text = p.read_text(encoding="utf-8")
        tree = ast.parse(text)

        imports: list[str] = []
        for node in tree.body:
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                seg = ast.get_source_segment(text, node)
                if seg:
                    imports.append(seg.strip())

        # de-dupe while preserving order
        seen = set()
        uniq: list[str] = []
        for line in imports:
            if line not in seen:
                seen.add(line)
                uniq.append(line)

        return "\n".join(uniq).strip() + ("\n" if uniq else "")
    except Exception:
        return ""


def _pack_payload(payload: Dict[str, Any], *, level: int = 6) -> str:
    lvl = int(level)
    if lvl < 1:
        lvl = 1
    elif lvl > 9:
        lvl = 9
    raw = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    comp = zlib.compress(raw, level=lvl)
    return _b64e(comp)


def parse_tagged_result(stdout_text: str, result_tag: str) -> Dict[str, Any]:
    """
    Extract the last tagged JSON payload printed by a cluster command.
    Expects lines like:  <<<RESULT>>>{...json...}
    """
    last = None
    for line in stdout_text.splitlines():
        if line.startswith(result_tag):
            last = line[len(result_tag) :]
    if last is None:
        raise ValueError(f"Result tag {result_tag!r} not found in output")
    try:
        return json.loads(last)
    except Exception as e:
        raise ValueError("Tagged result was not valid JSON") from e


class CommandError(RuntimeError):
    """
    Raised when the remote Databricks command returns ok=false.

    Attributes
    ----------
    error: str
        Short error message from remote.
    traceback: str
        Remote traceback (string) if provided.
    raw: dict
        Full parsed JSON payload from the tagged output line.
    """

    def __init__(self, error: str, traceback: str = "", raw: Optional[Dict[str, Any]] = None):
        self.error = error
        self.traceback = traceback
        self.raw = raw or {}
        msg = error if not traceback else f"{error}\n{traceback}"
        super().__init__(msg.rstrip())


class CommandResultParseError(ValueError):
    """
    Raised when the command output cannot be parsed or decoded.
    """

    def __init__(self, message: str, stdout_text: Optional[str] = None):
        self.stdout_text = stdout_text
        super().__init__(message)


@dataclass
class CallableSerdeMixin:
    """
    Encapsulates a callable + serialization strategy + Databricks command generation.

    - Same Python (major, minor): prefer dill
    - Different Python (major, minor): exec(imports + source) with captured globals/freevars
    """

    fn: Callable[..., Any]
    package_root: Tuple[str, ...] = field(default_factory=tuple)
    ALLOW_EXEC_SOURCE: bool = True

    @classmethod
    def from_callable(cls, fn: Callable[..., Any]) -> "CallableSerdeMixin":
        if isinstance(fn, CallableSerdeMixin):
            return fn
        return cls(fn=fn, package_root=_infer_package_root(fn))

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        return self.fn(*args, **kwargs)

    # ---------- pickle protocol ----------
    def __getstate__(self) -> Dict[str, Any]:
        if _is_python_function(self.fn):
            src = None
            try:
                src = textwrap.dedent(inspect.getsource(self.fn))
            except Exception:
                src = None

            payload: Dict[str, Any] = {
                "__callable__": True,
                "pyver": list(_pyver()),
                "name": getattr(self.fn, "__name__", None),
                "qualname": getattr(self.fn, "__qualname__", None),
                "module": getattr(self.fn, "__module__", None),
                "imports": _capture_module_imports(self.fn),  # NEW
                "source": src,
                "dill_b64": None,
                "env": _capture_exec_env(self.fn),
            }

            dumped = _safe_dill_dumps(self.fn)
            if dumped is not None:
                payload["dill_b64"] = _b64e(dumped)

            return {
                "fn": payload,
                "package_root": tuple(self.package_root),
                "ALLOW_EXEC_SOURCE": bool(self.ALLOW_EXEC_SOURCE),
            }

        dumped = _safe_dill_dumps(self.fn)
        if dumped is None:
            raise ValueError("Callable object could not be dill-serialized")
        return {
            "fn": {"__callable__": False, "dill_b64": _b64e(dumped)},
            "package_root": tuple(self.package_root),
            "ALLOW_EXEC_SOURCE": bool(self.ALLOW_EXEC_SOURCE),
        }

    def __setstate__(self, state: Dict[str, Any]) -> None:
        self.package_root = tuple(state.get("package_root") or ())
        self.ALLOW_EXEC_SOURCE = bool(state.get("ALLOW_EXEC_SOURCE", True))

        fn_payload = state["fn"]

        if isinstance(fn_payload, dict) and fn_payload.get("__callable__") is False:
            self.fn = _safe_dill_loads(_b64d(fn_payload["dill_b64"]))
            return

        if not isinstance(fn_payload, dict) or fn_payload.get("__callable__") is not True:
            raise ValueError("Invalid callable payload")

        src_pyver = tuple(fn_payload.get("pyver") or ())
        cur_pyver = _pyver()

        if src_pyver == cur_pyver and fn_payload.get("dill_b64"):
            try:
                self.fn = _safe_dill_loads(_b64d(fn_payload["dill_b64"]))
                if callable(self.fn):
                    return
            except Exception:
                pass

        if not self.ALLOW_EXEC_SOURCE:
            raise ValueError("Exec-based restore disabled and dill path unavailable/failed")

        imports = fn_payload.get("imports") or ""
        source = fn_payload.get("source")
        name = fn_payload.get("name")
        env = fn_payload.get("env") or {}

        if not source:
            raise ValueError("No source available for exec-based restore")

        ns: Dict[str, Any] = {}

        # preload captured names
        for bucket in ("globals", "freevars"):
            items = (env.get(bucket) or {})
            for k, b64 in items.items():
                try:
                    ns[k] = _safe_dill_loads(_b64d(b64))
                except Exception:
                    pass

        # NEW: run module imports first
        if imports.strip():
            exec(imports, ns, ns)

        exec(source, ns, ns)

        if name and name in ns and callable(ns[name]):
            self.fn = ns[name]
            return

        cands = [v for v in ns.values() if callable(v)]
        if cands:
            self.fn = cands[-1]
            return

        raise ValueError("exec(source) succeeded but no callable could be recovered")

    # ---------- Databricks command generation ----------
    def to_command(
        self,
        *,
        args: Tuple[Any, ...] = (),
        kwargs: Optional[Dict[str, Any]] = None,
        env_keys: Iterable[str] = (),
        env_variables: Optional[Dict[str, str]] = None,
        use_dill: bool,
        byte_limit: int = 0,
        result_tag: str = "<<<RESULT>>>",
        compress_input_payload: bool = True,  # NEW
        payload_compression_level: int = 6,  # NEW (1..9)
    ) -> str:
        if kwargs is None:
            kwargs = {}
        if env_variables is None:
            env_variables = {}

        # capture env vars from client process
        client_env: Dict[str, str] = {}
        for k in env_keys:
            v = os.environ.get(k)
            if v is not None:
                client_env[k] = v
        client_env.update(env_variables)

        if not _is_python_function(self.fn):
            raise ValueError("to_command supports Python functions/methods/lambdas only")

        src = None
        try:
            src = textwrap.dedent(inspect.getsource(self.fn))
        except Exception:
            src = None

        imports = _capture_module_imports(self.fn)

        callable_payload: Dict[str, Any] = {
            "__callable__": True,
            "pyver": list(_pyver()),
            "name": getattr(self.fn, "__name__", None),
            "qualname": getattr(self.fn, "__qualname__", None),
            "module": getattr(self.fn, "__module__", None),
            "imports": imports,
            "source": src,
            "dill_b64": None,
            "env": _capture_exec_env(self.fn),
        }
        dumped_fn = _safe_dill_dumps(self.fn)
        if dumped_fn is not None:
            callable_payload["dill_b64"] = _b64e(dumped_fn)

        # args/kwargs transport
        dumped_args = _safe_dill_dumps(args)
        dumped_kwargs = _safe_dill_dumps(kwargs)
        if dumped_args is None or dumped_kwargs is None:
            raise ValueError("Failed to dill-serialize args/kwargs")

        args_pack: Dict[str, Any] = {"kind": "dill", "b64": _b64e(dumped_args)}
        kwargs_pack: Dict[str, Any] = {"kind": "dill", "b64": _b64e(dumped_kwargs)}

        if not byte_limit:
            byte_limit = 512 * 1024

        payload = {
            "callable": callable_payload,
            "use_dill": bool(use_dill),
            "args": args_pack,
            "kwargs": kwargs_pack,
            "env": client_env,
            "result_tag": result_tag,
            "byte_limit": byte_limit,
        }

        # NEW: compress the embedded input payload to keep command size small
        if compress_input_payload:
            payload_b64 = _pack_payload(payload, level=payload_compression_level)
            payload_bootstrap = f"""
_payload_b64 = {payload_b64!r}

def _load_payload(b64: str) -> dict:
    raw = base64.b64decode(b64.encode("ascii"))
    raw = zlib.decompress(raw)
    return json.loads(raw.decode("utf-8"))

_payload = _load_payload(_payload_b64)
""".rstrip()
        else:
            payload_bootstrap = f"_payload = {payload!r}"

        return f"""
# --- generated by yggdrasil.ser.CallableSerdeMixin.to_command ---
import base64, json, os, traceback, zlib
import dill
from yggdrasil.databricks import *

{payload_bootstrap}

def _b64d(s: str) -> bytes:
    return base64.b64decode(s.encode("ascii"))

def _try_dill_load(b64: str):
    return dill.loads(_b64d(b64))

def _hydrate_env(env_dict):
    ns = {{}}
    for bucket in ("globals", "freevars"):
        items = (env_dict or {{}}).get(bucket, {{}}) or {{}}
        for name, b64 in items.items():
            try:
                ns[name] = _try_dill_load(b64)
            except Exception:
                pass
    return ns

def _load_args(pack):
    if pack["kind"] == "dill":
        return dill.loads(_b64d(pack["b64"]))
    return tuple(json.loads(pack["text"]))

def _load_kwargs(pack):
    if pack["kind"] == "dill":
        return dill.loads(_b64d(pack["b64"]))
    return dict(json.loads(pack["text"]))

def _load_callable(cpack, use_dill: bool):
    # Prefer dill when requested/available
    if use_dill and cpack.get("dill_b64"):
        fn = _try_dill_load(cpack["dill_b64"])
        if callable(fn):
            return fn

    src = cpack.get("source")
    name = cpack.get("name")
    imports = cpack.get("imports") or ""
    if not src:
        raise ValueError("No source available for exec-based restore")

    ns = {{}}
    ns.update(_hydrate_env(cpack.get("env") or {{}}))

    if imports.strip():
        exec(imports, ns, ns)

    exec(src, ns, ns)

    if name and name in ns and callable(ns[name]):
        return ns[name]

    cands = [v for v in ns.values() if callable(v)]
    if cands:
        return cands[-1]

    raise ValueError("exec(source) ran but no callable was recovered")

def _emit(tag: str, obj: dict):
    print(tag + json.dumps(obj, ensure_ascii=False, separators=(",", ":")))

def _zlib_level(n: int, limit: int) -> int:
    r = n / max(1, limit)
    if r >= 32: return 1
    if r >= 16: return 2
    if r >= 8:  return 3
    if r >= 4:  return 4
    if r >= 2:  return 5
    if r >= 1.25: return 6
    return 7

for k, v in (_payload.get("env") or {{}}).items():
    if v is not None:
        os.environ[str(k)] = str(v)

tag = _payload.get("result_tag", "<<<RESULT>>>")
byte_limit = int((_payload.get("byte_limit") or 2_000_000))

try:
    use_dill = bool(_payload.get("use_dill", False))
    fn = _load_callable(_payload["callable"], use_dill=use_dill)
    args = _load_args(_payload["args"])
    kwargs = _load_kwargs(_payload["kwargs"])

    out = fn(*args, **kwargs)
    out_raw = dill.dumps(out, recurse=True)

    if len(out_raw) > byte_limit:
        lvl = _zlib_level(len(out_raw), byte_limit)
        out_comp = zlib.compress(out_raw, level=lvl)
        out_b64 = base64.b64encode(out_comp).decode("ascii")
        _emit(tag, {{"ok": True, "encoding": f"dill+zlib{{lvl}}+b64", "payload": out_b64}})
    else:
        out_b64 = base64.b64encode(out_raw).decode("ascii")
        _emit(tag, {{"ok": True, "encoding": "dill+b64", "payload": out_b64}})

except Exception as e:
    tb = traceback.format_exc()
    err = {{"cls": e.__class__.__name__, "msg": str(e), "tb": tb}}
    _emit(tag, {{"ok": False, "err": err}})

# --- end generated command ---
""".lstrip()

    @staticmethod
    def parse_command_result(
        stdout_text: str,
        *,
        result_tag: str = "<<<RESULT>>>",
        decode: bool = True,
    ) -> Any:
        """
        Parse Databricks command stdout and return decoded result or raise.

        Protocol:
          prints lines like:  <<<RESULT>>>{...json...}
          last tagged line wins

        If ok=false -> raise DatabricksCommandError
        If ok=true and decode=True:
          - encoding == "dill+b64": dill.loads(base64(payload))
          - encoding == "repr": returns payload as-is
        If decode=False: returns the raw parsed dict.
        """
        last: Optional[str] = None
        for line in stdout_text.splitlines():
            if line.startswith(result_tag):
                last = line[len(result_tag):]
            else:
                print(line)

        if last is None:
            raise CommandResultParseError(
                f"Result tag {result_tag!r} not found in command output",
                stdout_text=stdout_text,
            )

        try:
            msg: Dict[str, Any] = json.loads(last)
        except Exception as e:
            raise CommandResultParseError(
                "Tagged result is not valid JSON",
                stdout_text=stdout_text,
            ) from e

        if not decode:
            return msg

        ok = bool(msg.get("ok", False))
        if not ok:
            error = msg.get("err", {})
            error_class = str(error.get("cls", "RuntimeError"))
            error_message = str(error.get("msg", "Remote execution failed"))
            error_traceback = str(error.get("tb", ""))

            base = CommandError(
                error=error_message,
                traceback=error_traceback,
                raw=msg,
            )

            if error_class == "ModuleNotFound":
                raise ModuleNotFoundError(error_message) from base

            raise base

        encoding = msg.get("encoding")
        payload = msg.get("payload")

        if encoding == "dill+b64":
            if not isinstance(payload, str):
                raise CommandResultParseError(
                    "Expected base64 string payload for dill+b64 encoding",
                    stdout_text=stdout_text,
                )
            try:
                data = base64.b64decode(payload.encode("ascii"))
                return dill.loads(data)
            except Exception as e:
                raise CommandResultParseError(
                    "Failed to decode dill+b64 payload",
                    stdout_text=stdout_text,
                ) from e

        if encoding == "repr":
            return payload

        if encoding.startswith("dill+zlib") and encoding.endswith("+b64"):
            if not isinstance(payload, str):
                raise CommandResultParseError("Expected base64 string payload for zlib encoding",
                                              stdout_text=stdout_text)
            try:
                data = base64.b64decode(payload.encode("ascii"))
                raw = zlib.decompress(data)
                return dill.loads(raw)
            except Exception as e:
                raise CommandResultParseError("Failed to decode dill+zlib+b64 payload", stdout_text=stdout_text) from e

        raise CommandResultParseError(
            f"Unknown result encoding: {encoding!r}",
            stdout_text=stdout_text,
        )