from __future__ import annotations

import functools
import hashlib
import logging
import os
import sys
import threading
from pathlib import Path
from typing import Any, Callable, Optional

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


_DEFAULT_SIZE_THRESHOLD: int = 50 * 1024 * 1024

def estimate_arg_size(obj: Any) -> int:
    try:
        if hasattr(obj, "_data") and hasattr(obj, "_filename"):
            return len(obj._data)

        if hasattr(obj, "nbytes"):
            return int(obj.nbytes)

        if hasattr(obj, "memory_usage"):
            usage = obj.memory_usage(deep=True)
            return int(usage.sum() if hasattr(usage, "sum") else usage)

        if isinstance(obj, (bytes, bytearray)):
            return len(obj)

        if isinstance(obj, str):
            return len(obj)

        if isinstance(obj, (list, tuple)):
            total = sys.getsizeof(obj)
            for item in obj:
                total += _shallow_size(item)
            return total

        if isinstance(obj, dict):
            total = sys.getsizeof(obj)
            for k, v in obj.items():
                total += _shallow_size(k) + _shallow_size(v)
            return total

        return sys.getsizeof(obj)
    except Exception:
        return 0


def _shallow_size(obj: Any) -> int:
    if hasattr(obj, "nbytes"):
        return int(obj.nbytes)
    if isinstance(obj, (bytes, bytearray, str)):
        return len(obj)
    return sys.getsizeof(obj)


class DataOffloader:
    def __init__(self, api_url: str, api_key: str, size_threshold: int | None = None) -> None:
        self.api_url = api_url
        self.api_key = api_key
        self._size_threshold = size_threshold if size_threshold is not None else _DEFAULT_SIZE_THRESHOLD
        self._pending_keys: list[str] = []
        self._keys_lock = threading.Lock()
        self._upload_cache: dict[str, str] = {}

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

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._keys_lock = threading.Lock()

    def offload_large_args(
        self,
        func: Callable[..., Any],
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> tuple[Callable[..., Any], tuple[Any, ...], dict[str, Any], list[str]]:
        arg_refs: dict[int, str] = {}
        kwarg_refs: dict[str, str] = {}
        data_keys: list[str] = []

        modified_args = list(args)
        modified_kwargs = dict(kwargs)

        for i, arg in enumerate(args):
            if estimate_arg_size(arg) < self._size_threshold:
                continue
            data_key = self._upload(arg)
            if data_key is None:
                continue
            arg_refs[i] = data_key
            data_keys.append(data_key)
            modified_args[i] = None

        for key, val in kwargs.items():
            if estimate_arg_size(val) < self._size_threshold:
                continue
            data_key = self._upload(val)
            if data_key is None:
                continue
            kwarg_refs[key] = data_key
            data_keys.append(data_key)
            modified_kwargs[key] = None

        if not data_keys:
            return func, args, kwargs, []

        logger.info("Offloaded %d arg(s) to backend data store", len(data_keys))

        _original_func = func
        _arg_refs = dict(arg_refs)
        _kwarg_refs = dict(kwarg_refs)
        _api_url = self.api_url

        @functools.wraps(func)
        def _resolving_wrapper(*call_args: Any, **call_kwargs: Any) -> Any:
            import cloudpickle as _cp
            import httpx as _hx

            resolved_args = list(call_args)
            resolved_kwargs = dict(call_kwargs)

            for idx, dk in _arg_refs.items():
                resp = _hx.get(f"{_api_url}/v1/data/{dk}", timeout=600.0)
                resp.raise_for_status()
                resolved_args[idx] = _cp.loads(resp.content)

            for kw, dk in _kwarg_refs.items():
                resp = _hx.get(f"{_api_url}/v1/data/{dk}", timeout=600.0)
                resp.raise_for_status()
                resolved_kwargs[kw] = _cp.loads(resp.content)

            return _original_func(*resolved_args, **resolved_kwargs)

        return _resolving_wrapper, tuple(modified_args), modified_kwargs, data_keys

    def _upload(self, arg: Any) -> Optional[str]:
        try:
            import cloudpickle

            serialized = cloudpickle.dumps(arg)
        except Exception as exc:
            logger.warning("Cannot serialise arg for offload: %s", exc)
            return None

        content_hash = hashlib.sha256(serialized).hexdigest()
        cached_key = self._upload_cache.get(content_hash)
        if cached_key is not None:
            logger.debug("Upload cache hit hash=%s → key=%s", content_hash[:12], cached_key)
            return cached_key

        try:
            import httpx

            resp = httpx.post(
                f"{self.api_url}/v1/data/upload",
                content=serialized,
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/octet-stream",
                },
                timeout=300.0,
            )
            if resp.status_code != 200:
                logger.warning("Data upload failed: %d %s", resp.status_code, resp.text[:200])
                return None

            data_key = resp.json().get("data_key")
            if data_key:
                self._upload_cache[content_hash] = data_key
                with self._keys_lock:
                    self._pending_keys.append(data_key)
            return data_key

        except Exception as exc:
            logger.warning("Data upload error: %s", exc)
            return None

    def flush_keys(self, keys: list[str]) -> None:
        if not keys:
            return
        try:
            import httpx
        except ImportError:
            return

        for dk in keys:
            try:
                httpx.delete(
                    f"{self.api_url}/v1/data/{dk}",
                    headers={"Authorization": f"Bearer {self.api_key}"},
                    timeout=10.0,
                )
            except Exception:
                pass

        with self._keys_lock:
            remaining = set(self._pending_keys) - set(keys)
            self._pending_keys = list(remaining)

        logger.debug("Flushed %d data key(s)", len(keys))

    def flush_all(self) -> None:
        with self._keys_lock:
            keys = list(self._pending_keys)
        if keys:
            logger.info("Flushing %d remaining data key(s)…", len(keys))
            self.flush_keys(keys)
            self._upload_cache.clear()


class _RemoteFile(os.PathLike):
    def __init__(self, data: bytes, filename: str) -> None:
        self._data = data
        self._filename = filename
        self._path: str | None = None

    def _materialize(self) -> str:
        if self._path is None:
            import tempfile

            tmpdir = tempfile.mkdtemp(prefix="verlex_files_")
            self._path = os.path.join(tmpdir, self._filename)
            with open(self._path, "wb") as fh:
                fh.write(self._data)
        return self._path

    def __fspath__(self) -> str:
        return self._materialize()

    def __str__(self) -> str:
        return self._materialize()

    def __repr__(self) -> str:
        return f"RemoteFile({self._filename!r}, {len(self._data)} bytes)"


class _RemoteDir(os.PathLike):
    def __init__(self, data: bytes, dirname: str, file_count: int) -> None:
        self._data = data
        self._dirname = dirname
        self._file_count = file_count
        self._path: str | None = None

    def _materialize(self) -> str:
        if self._path is None:
            import tempfile
            import zipfile
            import io

            tmpdir = tempfile.mkdtemp(prefix="verlex_dirs_")
            self._path = os.path.join(tmpdir, self._dirname)
            with zipfile.ZipFile(io.BytesIO(self._data), "r") as zf:
                zf.extractall(self._path)
        return self._path

    def __fspath__(self) -> str:
        return self._materialize()

    def __str__(self) -> str:
        return self._materialize()

    def __repr__(self) -> str:
        return f"RemoteDir({self._dirname!r}, {self._file_count} files, {len(self._data)} bytes)"


def _zip_directory(dir_path: Path) -> tuple[bytes, int]:
    import zipfile
    import io

    buf = io.BytesIO()
    file_count = 0
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for root, _dirs, files in os.walk(dir_path):
            for fname in files:
                full = os.path.join(root, fname)
                arcname = os.path.relpath(full, dir_path)
                zf.write(full, arcname)
                file_count += 1
    return buf.getvalue(), file_count


def _is_local_file(obj: Any) -> bool:
    if isinstance(obj, Path):
        return obj.is_file()
    if isinstance(obj, str) and len(obj) < 4096:
        if obj.startswith(("http://", "https://", "s3://", "gs://")):
            return False
        return os.path.isfile(obj)
    return False


def _is_local_dir(obj: Any) -> bool:
    if isinstance(obj, Path):
        return obj.is_dir()
    if isinstance(obj, str) and len(obj) < 4096:
        if obj.startswith(("http://", "https://", "s3://", "gs://")):
            return False
        return os.path.isdir(obj)
    return False


def _resolve_one(obj: Any) -> Any:
    if isinstance(obj, Path) and obj.is_file():
        return _RemoteFile(obj.read_bytes(), obj.name)
    if isinstance(obj, str) and os.path.isfile(obj):
        return _RemoteFile(
            Path(obj).read_bytes(),
            os.path.basename(obj),
        )
    if isinstance(obj, Path) and obj.is_dir():
        zip_bytes, file_count = _zip_directory(obj)
        return _RemoteDir(zip_bytes, obj.name, file_count)
    if isinstance(obj, str) and os.path.isdir(obj):
        zip_bytes, file_count = _zip_directory(Path(obj))
        return _RemoteDir(zip_bytes, os.path.basename(obj.rstrip("/\\")), file_count)
    return obj


def resolve_file_args(
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> tuple[tuple[Any, ...], dict[str, Any], int]:
    files_found = 0

    def _walk(obj: Any, depth: int = 0) -> Any:
        nonlocal files_found
        if _is_local_file(obj):
            files_found += 1
            return _resolve_one(obj)
        if _is_local_dir(obj):
            files_found += 1
            return _resolve_one(obj)
        if depth > 0:
            return obj
        if isinstance(obj, (list, tuple)):
            resolved = [_walk(item, depth + 1) for item in obj]
            return type(obj)(resolved)
        if isinstance(obj, dict):
            return {k: _walk(v, depth + 1) for k, v in obj.items()}
        return obj

    new_args = tuple(_walk(a) for a in args)
    new_kwargs = {k: _walk(v) for k, v in kwargs.items()}
    return new_args, new_kwargs, files_found
