from __future__ import annotations

import base64
import json
import os
import sys
import threading
import time
from dataclasses import dataclass
from typing import Any, Callable, TypeVar

import cloudpickle
import httpx

from verlex.data_ref import DataOffloader, resolve_file_args
from verlex.errors import (
    ExecutionError,
    InsufficientCreditsError,
    InvalidAPIKeyError,
    JobFailedError,
    JobTimeoutError,
    NetworkError,
    NotAuthenticatedError,
    RateLimitError,
    SerializationError,
)

T = TypeVar("T")

DEFAULT_API_URL = "https://api.verlex.dev"

_current_session: AuthSession | None = None


_DEFAULT_EARLY_PROGRESS: tuple[float, str] = (0.05, "Starting...")

_GATEWAY_MILESTONES: list[tuple[str, float, str]] = [
    ("Received",           0.12, "Received"),
    ("Analyzing",          0.15, "Analyzing..."),
    ("Optimizing",         0.20, "Optimizing..."),
    ("Selecting",          0.23, "Selecting resources..."),
    ("Launching",          0.28, "Launching..."),
    ("Resources allocated", 0.32, "Resources allocated"),
    ("Provisioning",       0.35, "Provisioning..."),
    ("Starting up",        0.40, "Starting up..."),
    ("Booting",            0.45, "Booting..."),
    ("Installing",         0.50, "Installing packages..."),
    ("Environment ready",  0.55, "Ready"),
    ("Deploying",          0.58, "Deploying code..."),
    ("Code ready",         0.60, "Code deployed"),
    ("Running function",   0.65, "Running function..."),
    ("Collecting result",  0.90, "Collecting result..."),
    ("completed",          0.95, "Finishing up..."),
]


class _ProgressBar:
    _INTERPOLATION_CAP = 0.89

    def __init__(self, verbose: bool) -> None:
        self._verbose = verbose
        self._rich_progress: Any = None
        self._task: Any = None
        self._last_pct = -1
        self._backend_progress = 0.0
        self._display_progress = 0.0
        self._last_backend_time = 0.0
        self._last_desc = ""
        self._last_desc_shown = ""
        if not verbose:
            return
        try:
            from rich.progress import (
                BarColumn,
                Progress,
                TaskProgressColumn,
                TextColumn,
                TimeElapsedColumn,
            )

            self._rich_progress = Progress(
                TextColumn("  "),
                BarColumn(bar_width=30),
                TaskProgressColumn(),
                TextColumn("{task.description}"),
                TimeElapsedColumn(),
                transient=True,
            )
        except ImportError:
            pass

    def start(self) -> None:
        self._last_backend_time = time.time()
        if self._rich_progress:
            self._rich_progress.start()
            self._task = self._rich_progress.add_task("Submitting...", total=100)

    def update(self, progress: float | None, message: str | None) -> None:
        if not self._verbose:
            return
        now = time.time()
        backend_val = progress if progress is not None else self._backend_progress
        prev_backend = self._backend_progress

        if backend_val > self._backend_progress:
            self._backend_progress = backend_val
            self._display_progress = backend_val
            self._last_backend_time = now
        else:
            elapsed = now - self._last_backend_time
            creep = elapsed * 0.005
            interpolated = self._backend_progress + creep
            cap = min(self._INTERPOLATION_CAP, self._backend_progress + 0.20)
            self._display_progress = min(interpolated, cap)

        if backend_val > prev_backend and message:
            self._last_desc = message
        desc = self._last_desc or message or ""

        pct = int(self._display_progress * 100)
        desc_changed = desc != self._last_desc_shown
        if pct == self._last_pct and not desc_changed:
            return
        self._last_pct = pct
        self._last_desc_shown = desc
        if self._rich_progress and self._task is not None:
            self._rich_progress.update(self._task, completed=pct, description=desc)
        else:
            bar_len = pct // 4
            bar = "\u2588" * bar_len + "\u2591" * (25 - bar_len)
            print(f"\r   [{bar}] {pct:3d}%  {desc:<30s}", end="", flush=True)

    def pause(self) -> None:
        if self._rich_progress:
            self._rich_progress.stop()
        elif self._verbose and self._last_pct >= 0:
            print("\r" + " " * 60 + "\r", end="", flush=True)

    def resume(self) -> None:
        if self._rich_progress:
            self._rich_progress.start()
        elif self._verbose and self._last_pct >= 0:
            pct = self._last_pct
            bar_len = pct // 4
            bar = "\u2588" * bar_len + "\u2591" * (25 - bar_len)
            print(f"\r   [{bar}] {pct:3d}%  {'':30s}", end="", flush=True)

    def finish(self) -> None:
        if self._last_pct < 0:
            return
        self._last_pct = -1
        if self._rich_progress:
            if self._task is not None:
                self._rich_progress.update(self._task, completed=100, description="Done")
            self._rich_progress.stop()
        elif self._verbose:
            print()


@dataclass
class AuthSession:
    api_key: str
    user_id: str | None = None
    email: str | None = None
    tier: str = "free"
    fast: bool = False


@dataclass
class JobResult:
    value: Any
    job_id: str
    execution_time_seconds: float = 0.0
    cost: float = 0.0

    def __repr__(self) -> str:
        return f"JobResult(value={self.value!r}, job_id='{self.job_id}', cost=${self.cost:.4f})"


class GateWay:
    def __init__(
        self,
        api_key: str | None = None,
        fast: bool = False,
        api_url: str | None = None,
        timeout: int | None = None,
        verbose: bool = True,
    ):
        self.api_key = (
            api_key
            or os.getenv("VERLEX_API_KEY")
            or os.getenv("GATEWAY_API_KEY")
            or (_current_session.api_key if _current_session else None)
        )
        self.fast = fast or (_current_session.fast if _current_session else False)
        self.api_url = api_url or os.getenv("VERLEX_API_URL", DEFAULT_API_URL)
        self.default_timeout = timeout
        self.verbose = verbose

        self._client: httpx.Client | None = None
        self._session: AuthSession | None = None
        self._active = False
        self._jobs: list[str] = []
        self._jobs_lock = threading.Lock()

        self._offloader: DataOffloader | None = None
        self._data_offload_threshold: int = 50 * 1024 * 1024
        self._warm_session_id: str | None = None
        self._header_offset = 0
        self._lines_above_bar = 0
        self._header_replaced = False

    def __getstate__(self):
        state = self.__dict__.copy()
        state.pop("_jobs_lock", None)
        state.pop("_client", None)
        state.pop("_session", None)
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._jobs_lock = threading.Lock()
        self._client = None
        self._session = None

    def __enter__(self) -> GateWay:
        if not self.api_key:
            raise NotAuthenticatedError(
                "No API key provided. Either pass api_key parameter or set VERLEX_API_KEY environment variable."
            )

        if self.verbose:
            print("Verlex - Initializing cloud execution...")
            self._lines_above_bar += 1

        self._client = httpx.Client(
            base_url=self.api_url,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
            timeout=httpx.Timeout(120.0, connect=30.0),
        )

        max_retries = 3
        last_exc: Exception | None = None
        for attempt in range(max_retries):
            try:
                response = self._client.get("/v1/auth/validate")
                if response.status_code == 401:
                    self._client.close()
                    raise InvalidAPIKeyError("Invalid API key")
                elif response.status_code == 200:
                    data = response.json()
                    self._session = AuthSession(
                        api_key=self.api_key,
                        user_id=data.get("user_id"),
                        email=data.get("email"),
                        tier=data.get("tier", "free"),
                        fast=self.fast,
                    )
                    self._data_offload_threshold = data.get(
                        "data_offload_threshold", 50 * 1024 * 1024
                    )
                    break
                else:
                    self._client.close()
                    raise NetworkError(
                        f"Unexpected response from Verlex API: {response.status_code} - {response.text}"
                    )
            except httpx.RequestError as e:
                last_exc = e
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)
                    continue
                self._client.close()
                raise NetworkError(f"Failed to connect to Verlex API: {e}") from e
            except Exception:
                self._client.close()
                raise

        self._active = True
        self._offloader = DataOffloader(
            api_url=self.api_url,
            api_key=self.api_key,
            size_threshold=self._data_offload_threshold,
        )

        if self.verbose:
            mode = "performance" if self.fast else "standard"
            print(f"   Ready | Mode: {mode}")
            self._header_offset += 1
            self._lines_above_bar += 1

        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        self._active = False

        if self._offloader:
            self._offloader.flush_all()
            self._offloader = None

        if self._client:
            self._client.close()
            self._client = None

        if self.verbose and self._jobs:
            print(f"\nVerlex session complete. Jobs: {len(self._jobs)}")

        return False

    def _ensure_active(self) -> None:
        if not self._active:
            raise RuntimeError(
                "GateWay must be used as a context manager:\n"
                "    with verlex.GateWay(api_key='...') as gw:\n"
                "        result = gw.run(my_function)"
            )

    def pre_warm(self, func: Callable | None = None) -> None:
        self._ensure_active()

        source_code: str | None = None
        if func is not None:
            import inspect
            import textwrap
            try:
                source_code = textwrap.dedent(inspect.getsource(func))
            except (OSError, TypeError):
                pass

        try:
            resp = self._client.post(
                "/v1/pre-warm",
                json={
                    "source_code": source_code,
                    "fast": self.fast,
                },
                timeout=10.0,
            )
            if resp.status_code == 200:
                data = resp.json()
                self._warm_session_id = data.get("warm_session_id")
        except Exception:
            pass

    def run(
        self,
        func: Callable[..., T],
        *args,
        gpu: str | None = None,
        gpu_count: int = 1,
        cpu: int | None = None,
        memory: str | None = None,
        timeout: int | None = None,
        **kwargs,
    ) -> T:
        self._ensure_active()

        args, kwargs, n_files = resolve_file_args(args, kwargs)
        if n_files and self.verbose:
            print(f"   {n_files} local file(s) will be transferred to the cloud")
            self._header_offset += 1
            self._lines_above_bar += 1

        data_keys: list[str] = []
        if self._offloader:
            func, args, kwargs, data_keys = self._offloader.offload_large_args(
                func, args, kwargs,
            )

        job_id = None
        try:
            job_id = self._submit_job(func, args, kwargs, gpu, gpu_count, cpu, memory, timeout)
            result = self._wait_for_job(job_id, timeout or self.default_timeout)
        except BaseException as exc:
            if job_id and not isinstance(exc, (JobFailedError, SerializationError, NetworkError)):
                if self.verbose:
                    print("\nCancelling remote job...")
                cancelled = self._cancel_job_remote(job_id)
                if self.verbose:
                    if cancelled:
                        print("Job cancelled on server.")
                    else:
                        print(f"Could not cancel job on server. Cancel manually: job {job_id}")
            raise
        finally:
            if data_keys and self._offloader:
                self._offloader.flush_keys(data_keys)

        return result

    _MODULE_TO_PACKAGE = {
        "cv2": "opencv-python",
        "PIL": "Pillow",
        "sklearn": "scikit-learn",
        "skimage": "scikit-image",
        "yaml": "PyYAML",
        "bs4": "beautifulsoup4",
        "dateutil": "python-dateutil",
        "dotenv": "python-dotenv",
        "jwt": "PyJWT",
        "serial": "pyserial",
        "usb": "pyusb",
        "magic": "python-magic",
        "Crypto": "pycryptodome",
    }

    _STDLIB: frozenset[str] = getattr(sys, "stdlib_module_names", None) or frozenset({
        "os", "sys", "re", "json", "math", "random", "time", "datetime",
        "collections", "itertools", "functools", "typing", "pathlib", "io",
        "pickle", "copy", "hashlib", "base64", "struct", "array", "threading",
        "multiprocessing", "subprocess", "socket", "ssl", "http", "urllib",
        "email", "html", "xml", "logging", "warnings", "unittest", "tempfile",
        "shutil", "glob", "fnmatch", "stat", "contextlib", "abc",
        "dataclasses", "enum", "types", "inspect", "traceback", "ast",
        "operator", "string", "textwrap", "numbers", "decimal", "fractions",
        "statistics", "cmath", "secrets", "argparse", "configparser", "csv",
        "sqlite3", "platform", "errno", "ctypes", "concurrent", "asyncio",
        "queue", "sched", "signal", "mmap", "codecs", "unicodedata",
        "pprint", "dis", "token", "tokenize", "builtins", "importlib",
        "pkgutil", "zipimport", "compileall", "py_compile", "zipfile",
        "tarfile", "gzip", "bz2", "lzma", "zlib", "webbrowser", "cgi",
        "uuid", "heapq", "bisect", "weakref", "pdb", "profile", "timeit",
        "__future__", "_thread",
    })

    @staticmethod
    def _get_func_imports(func, source_code: str | None) -> set[str]:
        import ast
        import types

        func_imports: set[str] = set()

        if source_code:
            try:
                tree = ast.parse(source_code)
                for node in ast.walk(tree):
                    if isinstance(node, ast.Import):
                        for alias in node.names:
                            func_imports.add(alias.name.split(".")[0])
                    elif isinstance(node, ast.ImportFrom):
                        if node.module:
                            func_imports.add(node.module.split(".")[0])
                    elif isinstance(node, ast.Call):
                        name = None
                        if isinstance(node.func, ast.Attribute) and node.func.attr == "import_module":
                            if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
                                name = node.args[0].value
                        elif isinstance(node.func, ast.Name) and node.func.id == "__import__":
                            if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
                                name = node.args[0].value
                        if name:
                            func_imports.add(name.split(".")[0])
            except SyntaxError:
                pass

        if hasattr(func, "__globals__"):
            for value in func.__globals__.values():
                if isinstance(value, types.ModuleType):
                    func_imports.add(value.__name__.split(".")[0])

        if hasattr(func, "__closure__") and func.__closure__:
            for cell in func.__closure__:
                try:
                    value = cell.cell_contents
                    if isinstance(value, types.ModuleType):
                        func_imports.add(value.__name__.split(".")[0])
                except ValueError:
                    pass

        return func_imports

    @staticmethod
    def _detect_local_packages(func, source_code: str | None) -> list[str]:
        func_imports = GateWay._get_func_imports(func, source_code)

        local_packages: dict[str, str] = {}
        try:
            import importlib.metadata
            for dist in importlib.metadata.distributions():
                name = dist.metadata["Name"]
                if name:
                    local_packages[name.lower()] = dist.version
        except Exception:
            pass

        pinned: list[str] = []
        seen: set[str] = set()
        for mod in sorted(func_imports):
            if mod in GateWay._STDLIB:
                continue
            pkg_name = GateWay._MODULE_TO_PACKAGE.get(mod, mod)
            pkg_lower = pkg_name.lower()
            if pkg_lower in seen:
                continue
            seen.add(pkg_lower)
            version = local_packages.get(pkg_lower)
            if version:
                pinned.append(f"{pkg_name}=={version}")

        return pinned

    @staticmethod
    def _detect_arg_packages(args: tuple, kwargs: dict) -> list[str]:
        modules: set[str] = set()
        _seen_types: set[type] = set()

        def _extract(obj: Any, depth: int = 0) -> None:
            obj_type = type(obj)
            if obj_type not in _seen_types:
                _seen_types.add(obj_type)
                mod = obj_type.__module__
                if mod and mod != "builtins":
                    modules.add(mod.split(".")[0])
            if depth < 2:
                if isinstance(obj, (list, tuple, set, frozenset)):
                    for item in obj:
                        _extract(item, depth + 1)
                elif isinstance(obj, dict):
                    for v in obj.values():
                        _extract(v, depth + 1)

        for a in args:
            _extract(a)
        for v in kwargs.values():
            _extract(v)

        modules -= GateWay._STDLIB
        modules.discard("verlex")
        if not modules:
            return []

        local_packages: dict[str, str] = {}
        try:
            import importlib.metadata
            for dist in importlib.metadata.distributions():
                name = dist.metadata["Name"]
                if name:
                    local_packages[name.lower()] = dist.version
        except Exception:
            return []

        pinned: list[str] = []
        seen: set[str] = set()
        for mod in sorted(modules):
            pkg_name = GateWay._MODULE_TO_PACKAGE.get(mod, mod)
            pkg_lower = pkg_name.lower()
            if pkg_lower in seen:
                continue
            seen.add(pkg_lower)
            version = local_packages.get(pkg_lower)
            if version:
                pinned.append(f"{pkg_name}=={version}")

        return pinned

    @staticmethod
    def _bundle_local_modules(
        func, source_code: str | None,
    ) -> tuple[str | None, list[str]]:
        import ast
        import importlib.metadata
        import importlib.util
        import io
        import tarfile

        func_imports = GateWay._get_func_imports(func, source_code)
        if not func_imports:
            return None, []

        pip_names: set[str] = set()
        local_packages: dict[str, str] = {}
        try:
            for dist in importlib.metadata.distributions():
                name = dist.metadata["Name"]
                if name:
                    pip_names.add(name.lower())
                    version = dist.metadata["Version"]
                    if version:
                        local_packages[name.lower()] = version
                    top_level = dist.read_text("top_level.txt")
                    if top_level:
                        for tl in top_level.strip().splitlines():
                            tl_clean = tl.strip().lower()
                            pip_names.add(tl_clean)
                            if tl_clean != name.lower():
                                local_packages[tl_clean] = f"{name}=={version}" if version else name
        except Exception:
            pass

        def _is_local_import(mod_name: str) -> bool:
            if mod_name in GateWay._STDLIB:
                return False
            pkg_name = GateWay._MODULE_TO_PACKAGE.get(mod_name, mod_name)
            if pkg_name.lower() in pip_names or mod_name.lower() in pip_names:
                return False
            return True

        def _resolve_local_module(mod_name: str) -> tuple[str, str] | None:
            try:
                spec = importlib.util.find_spec(mod_name)
                if not spec:
                    return None
                if spec.submodule_search_locations:
                    for loc in spec.submodule_search_locations:
                        loc = str(loc)
                        if "site-packages" in loc or "dist-packages" in loc:
                            continue
                        if os.path.isdir(loc):
                            return (mod_name, loc)
                elif spec.origin:
                    origin = str(spec.origin)
                    if "site-packages" in origin or "dist-packages" in origin:
                        return None
                    if os.path.isfile(origin):
                        return (mod_name, origin)
            except (ModuleNotFoundError, ValueError, ImportError):
                pass
            return None

        def _scan_imports_in_file(filepath: str) -> set[str]:
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    tree = ast.parse(f.read())
                imports: set[str] = set()
                for node in ast.walk(tree):
                    if isinstance(node, ast.Import):
                        for alias in node.names:
                            imports.add(alias.name.split(".")[0])
                    elif isinstance(node, ast.ImportFrom):
                        if node.module:
                            imports.add(node.module.split(".")[0])
                return imports
            except Exception:
                return set()

        def _scan_imports_in_dir(dirpath: str) -> set[str]:
            imports: set[str] = set()
            for root, _dirs, files in os.walk(dirpath):
                for fname in files:
                    if fname.endswith(".py"):
                        imports |= _scan_imports_in_file(os.path.join(root, fname))
            return imports

        visited: set[str] = set()
        to_visit: set[str] = set()
        local_modules: list[tuple[str, str]] = []
        all_scanned_imports: set[str] = set()

        for mod_name in func_imports:
            if _is_local_import(mod_name):
                to_visit.add(mod_name)

        while to_visit:
            mod_name = to_visit.pop()
            if mod_name in visited:
                continue
            visited.add(mod_name)

            resolved = _resolve_local_module(mod_name)
            if not resolved:
                continue
            local_modules.append(resolved)

            _, path = resolved
            if os.path.isdir(path) and not os.path.exists(os.path.join(path, "__init__.py")):
                import warnings
                warnings.warn(
                    f"Local module '{mod_name}' at '{path}' has no __init__.py "
                    f"and may not be importable in the cloud. "
                    f"Add an empty __init__.py to fix this.",
                    stacklevel=2,
                )

            _, path = resolved
            if os.path.isdir(path):
                child_imports = _scan_imports_in_dir(path)
            else:
                child_imports = _scan_imports_in_file(path)

            all_scanned_imports |= child_imports

            for child_mod in child_imports:
                if child_mod not in visited and _is_local_import(child_mod):
                    to_visit.add(child_mod)

        if not local_modules:
            return None, []

        extra_pip_deps: list[str] = []
        for mod_name in sorted(all_scanned_imports):
            if mod_name in GateWay._STDLIB:
                continue
            if _is_local_import(mod_name):
                continue
            pkg_name = GateWay._MODULE_TO_PACKAGE.get(mod_name, mod_name)
            pkg_lower = pkg_name.lower()
            version = local_packages.get(pkg_lower) or local_packages.get(mod_name.lower())
            if version and "==" in version:
                extra_pip_deps.append(version)
            elif version:
                extra_pip_deps.append(f"{pkg_name}=={version}")
            else:
                extra_pip_deps.append(pkg_name)

        buf = io.BytesIO()
        with tarfile.open(fileobj=buf, mode="w:gz") as tar:
            for mod_name, path in local_modules:
                if os.path.isdir(path):
                    def _exclude_pycache(tarinfo):
                        if "__pycache__" in tarinfo.name:
                            return None
                        return tarinfo
                    tar.add(path, arcname=mod_name, filter=_exclude_pycache)
                else:
                    tar.add(path, arcname=os.path.basename(path))

        archive_bytes = buf.getvalue()
        if len(archive_bytes) > 50 * 1024 * 1024:
            return None, extra_pip_deps

        return base64.b64encode(archive_bytes).decode("utf-8"), extra_pip_deps

    def _submit_job(
        self,
        func: Callable,
        args: tuple,
        kwargs: dict,
        gpu: str | None = None,
        gpu_count: int = 1,
        cpu: int | None = None,
        memory: str | None = None,
        timeout: int | None = None,
    ) -> str:
        try:
            payload = cloudpickle.dumps((func, args, kwargs))
            payload_b64 = base64.b64encode(payload).decode("utf-8")
        except Exception as e:
            raise SerializationError(f"Failed to serialize function: {e}") from e

        _MAX_PAYLOAD_BYTES = 50 * 1024 * 1024
        if len(payload) > _MAX_PAYLOAD_BYTES:
            raise SerializationError(
                f"Serialized function + arguments is "
                f"{len(payload) // 1024 // 1024} MB (limit: 50 MB). "
                f"Pass large arrays/DataFrames via verlex.upload() and reference "
                f"them by file path inside your function instead."
            )

        import ast
        import inspect
        import textwrap
        source_code: str | None = None
        try:
            source_code = textwrap.dedent(inspect.getsource(func))
        except (OSError, TypeError):
            pass

        file_source: str | None = None
        try:
            src_file = inspect.getfile(func)
            if src_file and not src_file.startswith("<"):
                with open(src_file, "r", encoding="utf-8") as _f:
                    file_source = _f.read()
        except (OSError, TypeError):
            pass

        combined_source = source_code or ""
        if file_source and file_source != combined_source:
            combined_source = file_source + "\n" + (source_code or "")

        pip_packages = self._detect_local_packages(func, combined_source or source_code)

        arg_packages = self._detect_arg_packages(args, kwargs)
        if arg_packages:
            pip_packages = sorted(set(pip_packages) | set(arg_packages))

        local_modules_b64, extra_pip_deps = self._bundle_local_modules(func, source_code)
        if extra_pip_deps:
            pip_packages = sorted(set(pip_packages) | set(extra_pip_deps))

        effective_timeout = timeout or self.default_timeout

        resources: dict[str, Any] = {}
        if gpu:
            resources["gpu"] = gpu
            resources["gpu_count"] = gpu_count
        if cpu:
            resources["cpu"] = cpu
        if memory:
            resources["memory"] = memory

        import sys
        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"

        request_data = {
            "function": {
                "payload": payload_b64,
                "type": "cloudpickle",
                "python_version": python_version,
                "source_code": source_code,
                "pip_packages": pip_packages,
                "local_modules": local_modules_b64,
            },
            "resources": resources,
            "estimate": {
                "estimated_duration_seconds": min(effective_timeout, 300) if effective_timeout else 300,
            },
            "fast": self.fast,
        }

        if self._warm_session_id:
            request_data["warm_session_id"] = self._warm_session_id

        try:
            response = self._client.post(
                "/v1/jobs/submit",
                json=request_data,
                timeout=60.0,
            )
        except httpx.RequestError as e:
            raise NetworkError(f"Failed to submit job: {e}") from e

        if response.status_code == 401:
            raise InvalidAPIKeyError("Invalid API key")
        elif response.status_code == 402:
            data = response.json()
            raise InsufficientCreditsError(
                required=data.get("required", 0),
                available=data.get("available", 0),
            )
        elif response.status_code == 429:
            raise RateLimitError(retry_after=int(response.headers.get("Retry-After", 60)))
        elif response.status_code != 200:
            raise ExecutionError(f"Job submission failed: {response.text}")

        data = response.json()
        job_id = data.get("job_id")

        if not job_id:
            raise ExecutionError("API response missing job_id")

        with self._jobs_lock:
            self._jobs.append(job_id)

        return job_id

    _STALE_THRESHOLD = 600

    def _wait_for_job(self, job_id: str, timeout: int | None) -> Any:
        start_time = time.time()
        poll_interval = 0.5
        disconnected = False
        retry_backoff = 2.0

        last_status: str | None = None
        last_progress: float | None = None
        activity_ts = [time.time()]
        code_running = [False]

        bar = _ProgressBar(self.verbose)
        bar.start()

        output_thread: threading.Thread | None = None
        output_stop = threading.Event()
        not_found_retries = 0

        try:
            while True:
                if timeout:
                    elapsed = time.time() - start_time
                    if elapsed > timeout:
                        raise JobTimeoutError(job_id, timeout)

                try:
                    response = self._client.get(f"/v1/jobs/{job_id}")
                    if disconnected:
                        disconnected = False
                        retry_backoff = 2.0
                        activity_ts[0] = time.time()
                        if self.verbose:
                            bar.pause()
                            print("   Reconnected — resuming job tracking")
                            bar.resume()
                except httpx.RequestError:
                    if not disconnected:
                        disconnected = True
                        if self.verbose:
                            bar.pause()
                            print("   Connection lost — job is still running in the cloud, retrying...")
                            bar.resume()
                    time.sleep(min(retry_backoff, 30.0))
                    retry_backoff = min(retry_backoff * 1.5, 30.0)
                    continue

                if response.status_code == 429:
                    retry_after = 1
                    try:
                        retry_after = int(response.json().get("retry_after", 1))
                    except Exception:
                        pass
                    time.sleep(max(retry_after, 1))
                    poll_interval = min(poll_interval * 2, 5.0)
                    continue

                if response.status_code >= 500:
                    if not disconnected:
                        disconnected = True
                        if self.verbose:
                            bar.pause()
                            print("   Server temporarily unavailable — job is still running, retrying...")
                            bar.resume()
                    time.sleep(min(retry_backoff, 30.0))
                    retry_backoff = min(retry_backoff * 1.5, 30.0)
                    continue

                if response.status_code == 402:
                    raise InsufficientCreditsError(required=0, available=0)

                if response.status_code == 404:
                    not_found_retries += 1
                    if not_found_retries <= 5:
                        time.sleep(min(not_found_retries * 2, 10))
                        continue
                    raise ExecutionError(
                        "Job not found — the server may have restarted. "
                        "Your job's cloud resources were cleaned up. Please retry."
                    )

                if response.status_code != 200:
                    raise ExecutionError(f"Failed to get job status: {response.text}")

                not_found_retries = 0
                data = response.json()
                status = data.get("status")

                progress = data.get("progress")
                progress_message = data.get("progress_message")
                if progress is not None:
                    bar.update(progress, progress_message)

                if progress and progress >= 0.6:
                    code_running[0] = True

                if status != last_status or progress != last_progress:
                    activity_ts[0] = time.time()
                    last_status = status
                    last_progress = progress

                if self.verbose and output_thread is None and status in ("provisioning", "running", "completed", "failed"):
                    output_thread = threading.Thread(
                        target=self._output_poll_loop,
                        args=(job_id, bar, output_stop, activity_ts, code_running),
                        daemon=True,
                    )
                    output_thread.start()

                if code_running[0]:
                    idle = time.time() - activity_ts[0]
                    if idle > self._STALE_THRESHOLD:
                        mins = int(idle // 60)
                        raise JobTimeoutError(job_id, mins * 60)

                if status == "completed":
                    output_stop.set()
                    if output_thread and output_thread.is_alive():
                        output_thread.join(timeout=5.0)

                    bar.update(1.0, "Done")
                    bar.finish()

                    result_b64 = data.get("result")
                    if result_b64 is None:
                        try:
                            result_response = self._client.get(f"/v1/jobs/{job_id}/result")
                        except httpx.RequestError as e:
                            raise NetworkError(f"Failed to fetch job result: {e}") from e
                        if result_response.status_code != 200:
                            raise ExecutionError(f"Failed to get result: {result_response.text}")
                        result_b64 = result_response.json().get("result")

                    if result_b64:
                        try:
                            result_bytes = base64.b64decode(result_b64)
                            return cloudpickle.loads(result_bytes)
                        except Exception as e:
                            raise SerializationError(f"Failed to deserialize result: {e}") from e
                    return None

                elif status == "failed":
                    output_stop.set()
                    if output_thread and output_thread.is_alive():
                        output_thread.join(timeout=3.0)
                    bar.finish()
                    error_msg = data.get("error", "Unknown error")
                    logs = data.get("logs", "")
                    raise JobFailedError(job_id, error_msg, logs)

                elif status in ("pending", "queued", "scheduled", "provisioning", "running", "completing"):
                    if status in ("pending", "queued", "scheduled", "provisioning"):
                        time.sleep(poll_interval)
                        poll_interval = min(poll_interval * 1.3, 10.0)
                    else:
                        time.sleep(poll_interval)
                        poll_interval = min(poll_interval * 1.1, 3.0)
                else:
                    raise ExecutionError(f"Unknown job status: {status}")
        except BaseException:
            output_stop.set()
            if output_thread and output_thread.is_alive():
                output_thread.join(timeout=1.0)
            bar.finish()
            raise

    def _output_poll_loop(
        self,
        job_id: str,
        bar: _ProgressBar,
        stop: threading.Event,
        activity_ts: list[float],
        code_running: list[bool],
    ):
        offset = 0
        interval = 0.5
        while not stop.is_set():
            prev = offset
            offset = self._poll_output(job_id, offset, bar, code_running)
            if offset != prev:
                activity_ts[0] = time.time()
                interval = 0.5
            else:
                interval = min(interval * 1.5, 8.0)
            wait = interval
            while wait > 0 and not stop.is_set():
                step = min(wait, 0.2)
                time.sleep(step)
                wait -= step
        self._poll_output(job_id, offset, bar, code_running)

    def _poll_output(self, job_id: str, offset: int, bar: _ProgressBar, code_running: list[bool]) -> int:
        try:
            resp = self._client.get(
                f"/v1/jobs/{job_id}/output",
                params={"offset": offset},
                timeout=5.0,
            )
            if resp.status_code != 200:
                return offset
            out = resp.json()
            entries = out.get("entries", [])
            for entry in entries:
                text = entry.get("text", "")
                if not text:
                    continue
                if text.startswith("[Output] "):
                    bar.pause()
                    _out_text = text[9:]
                    print(_out_text, end="")
                    self._lines_above_bar += _out_text.count("\n")
                    bar.resume()
                elif text.startswith("[GateWay] "):
                    msg = text[10:].rstrip("\n")
                    for keyword, prog, label in _GATEWAY_MILESTONES:
                        if keyword.lower() in msg.lower():
                            bar.update(prog, label)
                            if prog >= 0.6:
                                code_running[0] = True
                            break
                    _header_detail = None
                    _is_confirmed = False
                    if msg.startswith("Running on: "):
                        _header_detail = msg[len("Running on: "):]
                        _is_confirmed = True
                    elif msg.startswith("Resources allocated: "):
                        _detail = msg[len("Resources allocated: "):]
                        at_idx = _detail.find(" @ $")
                        if at_idx != -1:
                            _detail = _detail[:at_idx]
                        _header_detail = _detail
                    if _header_detail and (_is_confirmed or not self._header_replaced):
                        self._header_replaced = True
                        bar.pause()
                        n = self._lines_above_bar
                        try:
                            sys.stdout.write(
                                f"\033[{n}A\rVerlex - {_header_detail}\033[K\033[{n}B\r"
                            )
                            sys.stdout.flush()
                        except Exception:
                            pass
                        bar.resume()
            return out.get("current_index", offset)
        except Exception:
            return offset

    def _cancel_job_remote(self, job_id: str, max_attempts: int = 10) -> bool:
        for attempt in range(1, max_attempts + 1):
            try:
                resp = self._client.post(
                    f"/v1/jobs/{job_id}/cancel",
                    timeout=10.0,
                )
                if resp.status_code == 200:
                    return True
                return False
            except (httpx.RequestError, httpx.TimeoutException):
                if attempt < max_attempts:
                    time.sleep(2.0)
                continue
        return False

    def run_async(self, func: Callable[..., T], *args, **kwargs) -> AsyncJob[T]:
        self._ensure_active()
        return AsyncJob(self, func, args, kwargs)

    @property
    def jobs(self) -> list[str]:
        return self._jobs.copy()


class AsyncJob:
    def __init__(
        self,
        gateway: GateWay,
        func: Callable,
        args: tuple,
        kwargs: dict,
    ):
        self._gateway = gateway
        self._func = func
        self._args = args
        self._kwargs = kwargs
        self._result: Any | None = None
        self._error: Exception | None = None
        self._completed = False
        self._started = False
        self._thread: threading.Thread | None = None
        self._job_id: str | None = None

    def start(self) -> AsyncJob:
        if self._started:
            return self

        self._started = True
        self._thread = threading.Thread(target=self._execute, daemon=True)
        self._thread.start()
        return self

    def _execute(self) -> None:
        data_keys: list[str] = []
        try:
            args, kwargs, _ = resolve_file_args(self._args, self._kwargs)

            func = self._func
            if self._gateway._offloader:
                func, args, kwargs, data_keys = self._gateway._offloader.offload_large_args(
                    func, args, kwargs,
                )

            self._job_id = self._gateway._submit_job(func, args, kwargs)

            self._result = self._gateway._wait_for_job(
                self._job_id,
                self._gateway.default_timeout,
            )
        except Exception as e:
            self._error = e
        finally:
            if data_keys and self._gateway._offloader:
                self._gateway._offloader.flush_keys(data_keys)
            self._completed = True

    def result(self, timeout: float | None = None) -> Any:
        if not self._started:
            self.start()

        if self._thread:
            self._thread.join(timeout=timeout)

        if not self._completed:
            raise TimeoutError(f"Job did not complete within {timeout} seconds")

        if self._error:
            raise self._error

        return self._result

    @property
    def completed(self) -> bool:
        return self._completed

    @property
    def job_id(self) -> str | None:
        return self._job_id

    def cancel(self) -> bool:
        if self._completed:
            return False

        self._error = ExecutionError("Job was cancelled")
        self._completed = True

        if self._job_id:
            self._gateway._cancel_job_remote(self._job_id)

        return True

    def __del__(self):
        self._completed = True


def auth(api_key: str | None = None, fast: bool = False) -> AuthSession:
    global _current_session

    api_key = api_key or os.getenv("VERLEX_API_KEY")
    if not api_key:
        raise NotAuthenticatedError("No API key provided")

    _current_session = AuthSession(
        api_key=api_key,
        fast=fast,
    )
    return _current_session


def get_session() -> AuthSession | None:
    return _current_session


def logout() -> None:
    global _current_session
    _current_session = None
    print("Logged out.")


def whoami() -> dict[str, Any] | None:
    if not _current_session:
        print("Not authenticated. Run verlex.auth() first.")
        return None

    api_url = os.getenv("VERLEX_API_URL", DEFAULT_API_URL)

    try:
        response = httpx.get(
            f"{api_url}/v1/auth/validate",
            headers={"Authorization": f"Bearer {_current_session.api_key}"},
            timeout=10.0,
        )
        if response.status_code == 200:
            data = response.json()
            print(f"Authenticated as: {data.get('email', 'user')}")
            print(f"Tier: {data.get('tier', 'free')}")
            return data
        else:
            print("Failed to get user info")
            return None
    except httpx.RequestError as e:
        print(f"Network error: {e}")
        return None


def cloud(
    func: Callable[..., T],
    *args,
    api_key: str | None = None,
    fast: bool = False,
    verbose: bool = True,
    gpu: str | None = None,
    cpu: int | None = None,
    memory: str | None = None,
    timeout: int | None = None,
    **kwargs,
) -> T:
    api_key = api_key or (_current_session.api_key if _current_session else None)
    fast = fast or (_current_session.fast if _current_session else False)
    with GateWay(api_key=api_key, fast=fast, verbose=verbose) as gw:
        return gw.run(func, *args, gpu=gpu, cpu=cpu, memory=memory, timeout=timeout, **kwargs)
