from __future__ import annotations

import atexit
import functools
import logging
import os
import sys
import threading
import time
from typing import Any, Callable, Dict, List, Optional, Set, Tuple

logger = logging.getLogger("verlex.overflow")

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


class _WrapRecord:
    __slots__ = ("original", "name", "namespace", "wrapper")

    def __init__(
        self,
        original: Callable[..., Any],
        name: str,
        namespace: dict[str, Any],
        wrapper: Callable[..., Any],
    ) -> None:
        self.original = original
        self.name = name
        self.namespace = namespace
        self.wrapper = wrapper


def _compute_exclude_prefixes() -> Tuple[str, ...]:
    import sysconfig as _sc

    paths: set[str] = set()

    for key in ("stdlib", "platstdlib", "purelib", "platlib"):
        p = _sc.get_paths().get(key, "")
        if p:
            paths.add(os.path.abspath(p))

    try:
        import site

        for sp in site.getsitepackages():
            paths.add(os.path.abspath(sp))
        usp = site.getusersitepackages()
        if isinstance(usp, str):
            paths.add(os.path.abspath(usp))
    except Exception:
        pass

    paths.add(os.path.dirname(os.path.abspath(__file__)))

    return tuple(sorted(paths, key=len, reverse=True))


class _OverflowMonitor:
    def __init__(
        self,
        fast: bool,
        threshold: float,
        check_interval: float,
        api_key: Optional[str],
        api_url: Optional[str],
        gpu_threshold: Optional[float] = None,
        verbose: bool = True,
        priority: bool | None = None,
        flexible: bool | None = None,
    ) -> None:
        self.fast = fast
        self.threshold = threshold
        self.check_interval = check_interval
        self.gpu_threshold = gpu_threshold if gpu_threshold is not None else threshold
        self.verbose = verbose
        self.api_key = (
            api_key
            or os.environ.get("VERLEX_API_KEY")
            or os.environ.get("GATEWAY_API_KEY")
        )
        self.api_url = api_url or os.environ.get(
            "VERLEX_API_URL", _DEFAULT_API_URL
        )

        self._running = False
        self._overloaded = False
        self._cpu_usage = 0.0
        self._mem_usage = 0.0
        self._gpu_util = 0.0
        self._gpu_mem_util = 0.0
        self._gpu_available = False
        self._nvml_handle: Any = None

        self._call_stack: list[dict[str, Any]] = []
        self._wrapped_functions: Dict[int, _WrapRecord] = {}
        self._heavy_signatures: Set[Tuple[str, str, int]] = set()
        self._scope_wrapped = False

        self._filename_cache: Dict[str, bool] = {}

        self._lock = threading.Lock()
        self._monitor_thread: Optional[threading.Thread] = None
        self._original_profile: Any = None

        self._gateway: Any = None
        self._gateway_lock = threading.Lock()
        self._gateway_failed = False

        self._warm_result: dict[str, Any] = {}

        self._exclude_prefixes = _compute_exclude_prefixes()

    def __getstate__(self):
        state = self.__dict__.copy()
        state.pop("_lock", None)
        state.pop("_gateway_lock", None)
        state.pop("_monitor_thread", None)
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._lock = threading.Lock()
        self._gateway_lock = threading.Lock()
        self._monitor_thread = None

    def start(self) -> None:
        try:
            import psutil  # noqa: F401
        except ImportError:
            raise ImportError(
                "psutil is required for overflow monitoring. "
                "Install it with:  pip install 'verlex[overflow]'"
            )

        if self._running:
            return

        self._init_gpu()

        self._running = True
        self._original_profile = sys.getprofile()
        sys.setprofile(self._profile_hook)

        self._monitor_thread = threading.Thread(
            target=self._resource_loop,
            daemon=True,
            name="verlex-overflow",
        )
        self._monitor_thread.start()

        atexit.register(self.stop)
        if self._gpu_available:
            logger.info(
                "Overflow monitor started  threshold=%.0f%%  gpu_threshold=%.0f%%  fast=%s",
                self.threshold,
                self.gpu_threshold,
                self.fast,
            )
        else:
            logger.info(
                "Overflow monitor started  threshold=%.0f%%  fast=%s  (GPU monitoring: off)",
                self.threshold,
                self.fast,
            )

    def _init_gpu(self) -> None:
        try:
            import pynvml

            pynvml.nvmlInit()
            device_count = pynvml.nvmlDeviceGetCount()
            if device_count == 0:
                logger.info("pynvml: no NVIDIA GPUs found — GPU monitoring disabled")
                return
            self._nvml_handle = pynvml.nvmlDeviceGetHandleByIndex(0)
            self._gpu_available = True
            name = pynvml.nvmlDeviceGetName(self._nvml_handle)
            if isinstance(name, bytes):
                name = name.decode("utf-8")
            logger.info(
                "GPU monitoring active  device=%s  gpu_threshold=%.0f%%",
                name,
                self.gpu_threshold,
            )
        except ImportError:
            logger.info("pynvml not installed — GPU monitoring disabled")
        except Exception as exc:
            logger.info("GPU monitoring unavailable: %s", exc)

    def stop(self) -> None:
        if not self._running:
            return

        self._running = False

        self._monitor_thread = None

        try:
            sys.setprofile(self._original_profile)
        except Exception:
            pass

        self._unwrap_all()

        if self._gateway is not None:
            try:
                self._gateway.verbose = False
                self._gateway.__exit__(None, None, None)
            except Exception:
                pass
            self._gateway = None

        if self._gpu_available:
            try:
                import pynvml

                pynvml.nvmlShutdown()
            except Exception:
                pass
            self._gpu_available = False
            self._nvml_handle = None

        logger.info("Overflow monitor stopped.")

    def _resource_loop(self) -> None:
        import psutil

        psutil.cpu_percent(interval=None)

        while self._running:
            try:
                self._cpu_usage = psutil.cpu_percent(interval=self.check_interval)
                self._mem_usage = psutil.virtual_memory().percent

                if self._gpu_available and self._nvml_handle is not None:
                    try:
                        import pynvml

                        util = pynvml.nvmlDeviceGetUtilizationRates(self._nvml_handle)
                        mem_info = pynvml.nvmlDeviceGetMemoryInfo(self._nvml_handle)
                        self._gpu_util = float(util.gpu)
                        self._gpu_mem_util = (
                            (mem_info.used / mem_info.total) * 100.0
                            if mem_info.total > 0
                            else 0.0
                        )
                    except Exception as exc:
                        logger.debug("GPU sampling failed: %s", exc)

                was_overloaded = self._overloaded

                cpu_mem_overloaded = (
                    self._cpu_usage > self.threshold
                    or self._mem_usage > self.threshold
                )

                gpu_overloaded = False
                if self._gpu_available:
                    gpu_overloaded = (
                        self._gpu_util > self.gpu_threshold
                        or self._gpu_mem_util > self.gpu_threshold
                    )

                self._overloaded = cpu_mem_overloaded or gpu_overloaded

                if self._overloaded and not was_overloaded:
                    parts = [
                        f"CPU={self._cpu_usage:.1f}%",
                        f"Memory={self._mem_usage:.1f}%",
                    ]
                    if self._gpu_available:
                        parts.append(f"GPU={self._gpu_util:.1f}%")
                        parts.append(f"GPU-Mem={self._gpu_mem_util:.1f}%")
                    logger.warning(
                        "System overloaded  %s  (threshold=%.0f%%)",
                        "  ".join(parts),
                        self.threshold,
                    )
                elif not self._overloaded and was_overloaded:
                    parts = [
                        f"CPU={self._cpu_usage:.1f}%",
                        f"Memory={self._mem_usage:.1f}%",
                    ]
                    if self._gpu_available:
                        parts.append(f"GPU={self._gpu_util:.1f}%")
                        parts.append(f"GPU-Mem={self._gpu_mem_util:.1f}%")
                    logger.info(
                        "Resources normalised  %s",
                        "  ".join(parts),
                    )
                    self._scope_wrapped = False

            except Exception as exc:
                logger.error("Resource sampling failed: %s", exc)
                time.sleep(self.check_interval)

    def _profile_hook(self, frame: Any, event: str, arg: Any) -> None:  # noqa: C901
        if not self._running:
            return

        filename = frame.f_code.co_filename

        if not self._is_user_code(filename):
            return

        if event == "call":
            self._call_stack.append(
                {
                    "name": frame.f_code.co_name,
                    "filename": filename,
                    "lineno": frame.f_code.co_firstlineno,
                    "start_cpu": self._cpu_usage,
                    "start_mem": self._mem_usage,
                    "start_gpu": self._gpu_util,
                    "start_gpu_mem": self._gpu_mem_util,
                }
            )

            if self._overloaded:
                self._try_wrap_from_frame(frame)
                if not self._scope_wrapped and frame.f_back is not None:
                    self._wrap_scope(frame.f_back)
                    self._scope_wrapped = True

        elif event == "return":
            if not self._call_stack:
                return
            info = self._call_stack.pop()

            started_ok = (
                info["start_cpu"] <= self.threshold
                and info["start_mem"] <= self.threshold
            )
            if self._gpu_available:
                started_ok = started_ok and (
                    info["start_gpu"] <= self.gpu_threshold
                    and info["start_gpu_mem"] <= self.gpu_threshold
                )
            if started_ok and self._overloaded:
                sig = (info["name"], info["filename"], info["lineno"])
                self._heavy_signatures.add(sig)
                self._try_wrap_from_frame(frame)

                if not self._scope_wrapped and frame.f_back is not None:
                    self._wrap_scope(frame.f_back)
                    self._scope_wrapped = True

                logger.warning(
                    "Heavy function detected: %s() at %s:%d — wrapped for cloud offload",
                    info["name"],
                    info["filename"],
                    info["lineno"],
                )

    def _is_user_code(self, filename: str) -> bool:
        cached = self._filename_cache.get(filename)
        if cached is not None:
            return cached
        result = self._check_user_code(filename)
        self._filename_cache[filename] = result
        return result

    def _check_user_code(self, filename: str) -> bool:
        if not filename or filename.startswith("<"):
            return False
        abs_path = os.path.abspath(filename)
        for prefix in self._exclude_prefixes:
            if abs_path.startswith(prefix):
                return False
        return True

    def _try_wrap_from_frame(self, frame: Any) -> None:
        func_name = frame.f_code.co_name
        code_obj = frame.f_code
        search = frame.f_back
        while search is not None:
            obj = search.f_globals.get(func_name)
            if (
                obj is not None
                and callable(obj)
                and not getattr(obj, "_overflow_wrapped", False)
                and hasattr(obj, "__code__")
                and obj.__code__ is code_obj
            ):
                self._register_wrap(obj, func_name, search.f_globals)
                return
            search = search.f_back

    def _wrap_scope(self, frame: Any) -> None:
        if frame is None:
            return
        ns = frame.f_globals
        for name, obj in list(ns.items()):
            if (
                callable(obj)
                and hasattr(obj, "__code__")
                and not getattr(obj, "_overflow_wrapped", False)
                and not name.startswith("_")
                and self._is_user_code(obj.__code__.co_filename)
            ):
                self._register_wrap(obj, name, ns)

    def _register_wrap(
        self,
        func_obj: Callable[..., Any],
        func_name: str,
        namespace: dict[str, Any],
    ) -> None:
        with self._lock:
            fid = id(func_obj)
            if fid in self._wrapped_functions:
                return

            wrapper = self._make_wrapper(func_obj)
            namespace[func_name] = wrapper

            self._wrapped_functions[fid] = _WrapRecord(
                original=func_obj,
                name=func_name,
                namespace=namespace,
                wrapper=wrapper,
            )

    def _make_wrapper(self, original: Callable[..., Any]) -> Callable[..., Any]:
        monitor = self

        @functools.wraps(original)
        def _overflow_wrapper(*args: Any, **kwargs: Any) -> Any:
            if monitor._overloaded and not monitor._gateway_failed:
                parts = [f"CPU={monitor._cpu_usage:.0f}%", f"Mem={monitor._mem_usage:.0f}%"]
                if monitor._gpu_available:
                    parts.append(f"GPU={monitor._gpu_util:.0f}%")
                    parts.append(f"GPU-Mem={monitor._gpu_mem_util:.0f}%")
                reason = "  ".join(parts)

                logger.info(
                    "Offloading %s() to cloud  (%s)",
                    original.__name__,
                    reason,
                )

                if monitor.verbose:
                    print()
                    print(f"  \u2191 Overflow \u2192 Cloud: {original.__name__}()")
                    print(f"    Reason: {reason}")
                    print()

                try:
                    result = monitor._cloud_execute(original, args, kwargs)
                    logger.info(
                        "%s() completed in cloud — returning result to local",
                        original.__name__,
                    )
                    if monitor.verbose:
                        print(f"  \u2713 {original.__name__}() completed in cloud")
                        print()
                    return result
                except Exception as exc:
                    logger.warning(
                        "Cloud offload of %s() failed (%s) — falling back to local",
                        original.__name__,
                        exc,
                    )
                    if monitor.verbose:
                        print(f"  \u26a0 Cloud failed — running {original.__name__}() locally")
                    return original(*args, **kwargs)
            return original(*args, **kwargs)

        _overflow_wrapper._overflow_wrapped = True  # type: ignore[attr-defined]
        _overflow_wrapper._original = original  # type: ignore[attr-defined]
        return _overflow_wrapper

    def _start_pre_warm(self, source_file: Optional[str] = None) -> None:
        source_code: str | None = None
        if source_file:
            try:
                with open(source_file, "r", encoding="utf-8") as f:
                    source_code = f.read()
            except Exception:
                pass

        def _do_warm():
            try:
                import httpx

                api_url = self.api_url or _DEFAULT_API_URL
                api_key = self.api_key
                if not api_key:
                    return
                resp = httpx.post(
                    f"{api_url}/v1/pre-warm",
                    json={
                        "source_code": source_code,
                        "fast": self.fast,
                    },
                    headers={"Authorization": f"Bearer {api_key}"},
                    timeout=10.0,
                )
                if resp.status_code == 200:
                    self._warm_result = resp.json()
            except Exception:
                pass

        threading.Thread(target=_do_warm, daemon=True).start()

    def _get_gateway(self) -> Any:
        if self._gateway is None and not self._gateway_failed:
            with self._gateway_lock:
                if self._gateway is None and not self._gateway_failed:
                    from verlex.client import GateWay

                    try:
                        gw = GateWay(
                            api_key=self.api_key,
                            fast=self.fast,
                            api_url=self.api_url,
                            verbose=self.verbose,
                        )
                        gw.__enter__()
                        self._gateway = gw
                    except Exception:
                        self._gateway_failed = True
                        raise
        return self._gateway

    def _cloud_execute(
        self,
        func: Callable[..., Any],
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> Any:
        gw = self._get_gateway()
        warm_id = self._warm_result.get("warm_session_id")
        if warm_id:
            gw._warm_session_id = warm_id
        return gw.run(func, *args, **kwargs)

    def _unwrap_all(self) -> None:
        with self._lock:
            for record in self._wrapped_functions.values():
                try:
                    record.namespace[record.name] = record.original
                except Exception as exc:
                    logger.error("Failed to unwrap %s: %s", record.name, exc)
            self._wrapped_functions.clear()
            self._heavy_signatures.clear()
            self._scope_wrapped = False


_active_monitor: Optional[_OverflowMonitor] = None


def overflow(
    fast: bool = False,
    threshold: float = 85.0,
    check_interval: float = 0.5,
    api_key: Optional[str] = None,
    api_url: Optional[str] = None,
    gpu_threshold: Optional[float] = None,
    verbose: bool = True,
    priority: bool | None = None,
    flexible: bool | None = None,
    _caller_frame: object = None,
) -> _OverflowMonitor:
    global _active_monitor

    if _active_monitor is not None:
        _active_monitor.stop()

    _active_monitor = _OverflowMonitor(
        fast=fast,
        threshold=threshold,
        check_interval=check_interval,
        api_key=api_key,
        api_url=api_url,
        gpu_threshold=gpu_threshold,
        verbose=verbose,
    )
    _active_monitor.start()

    caller_frame = _caller_frame
    if caller_frame is None:
        try:
            caller_frame = sys._getframe(1)
        except (ValueError, AttributeError):
            pass
    if caller_frame is not None:
        _active_monitor._wrap_scope(caller_frame)
        source_file = caller_frame.f_code.co_filename
        if source_file and not source_file.startswith("<"):
            _active_monitor._start_pre_warm(source_file)

    return _active_monitor
