"""Large namespace values moved to disk behind a handle that pages like job output.

A kernel past its memory cap after a cell moves each large top-level value to
a file in the output directory and binds the name to a `Paged` handle: a str,
bytes or list of lines is written as text or raw bytes, anything else is
pickled. The handle is small and pickles, so it survives a release.
"""

from __future__ import annotations

import ctypes
import gc
import itertools
import os
import sys
import tempfile
import types
from collections.abc import Iterator, Sized
from typing import Any, BinaryIO, cast

import albedo_api
import albedo_capture
import albedo_memory
import albedo_output
import albedo_state

# A value this large is moved: the size the saved state would skip anyway.
MOVE_MIN = albedo_state.STATE_MAX_VALUE
WRITE_CHUNK = 1 << 20
SAMPLE = 200


def size_of(value: object) -> int:
    """What a value holds, cheaply: exact for str, bytes and arrays, sampled
    for containers, shallow for anything else."""
    if isinstance(value, (str, bytes, bytearray)):
        return sys.getsizeof(value)
    if isinstance(value, memoryview):
        return value.nbytes
    nbytes = getattr(value, "nbytes", None)
    if isinstance(nbytes, int):
        return nbytes
    if isinstance(value, dict):
        items = list(itertools.islice(value.items(), SAMPLE))
        each = sum(sys.getsizeof(k) + sys.getsizeof(v) for k, v in items)
    elif isinstance(value, (list, tuple, set, frozenset)):
        items = list(itertools.islice(value, SAMPLE))
        each = sum(sys.getsizeof(item) + 8 for item in items)
    else:
        return sys.getsizeof(value)
    shallow = sys.getsizeof(value)
    return shallow + (each * len(value) // len(items) if items else 0)


def mebibytes(size: int) -> str:
    return albedo_memory.mebibytes(size)


class Paged:
    """A value held on disk. `read(offset, limit)` pages its bytes as text,
    `head()` and `tail()` preview it, `load()` brings it back whole (refused
    past the memory cap), and a text or line value iterates its lines."""

    def __init__(
        self, path: str, kind: str, size: int, length: int | None, container: str = ""
    ) -> None:
        self.path = path
        self.kind = kind  # "text", "bytes", "lines" or "pickle"
        self.size = size
        self.length = length
        self.container = container  # "list" or "tuple" for kind "lines"

    def __len__(self) -> int:
        if self.length is None:
            raise TypeError("the value on disk has no length")
        return self.length

    def __repr__(self) -> str:
        count = {"text": "characters", "bytes": "bytes", "lines": "lines"}.get(
            self.kind, "items"
        )
        held = f", {self.length:,} {count}" if self.length is not None else ""
        return (
            f"<on disk: {self.kind} ({mebibytes(self.size)}{held}) at {self.path};"
            " page it with .read(offset, limit), .head(), .tail(lines=)>"
        )

    def _require_text(self) -> None:
        if self.kind == "pickle":
            raise TypeError(
                f"a pickled {self.container or 'value'} cannot be paged;"
                " .load() brings it back if it fits under the memory cap"
            )

    def _bytes(self, offset: int, limit: int) -> bytes:
        try:
            with open(self.path, "rb") as file:
                file.seek(offset)
                return file.read(limit)
        except FileNotFoundError:
            raise FileNotFoundError(
                f"{self.path} is gone (spill files are kept for 14 days)"
            ) from None

    def read(self, offset: int = 0, limit: int = albedo_api.OUTPUT_PREVIEW) -> Any:
        """A page of at most `limit` bytes from byte `offset`, cut on
        characters; advance by the bytes the page used (`len(page.encode())`)."""
        self._require_text()
        if offset < 0 or not 4 <= limit <= albedo_api.OUTPUT_PREVIEW:
            raise ValueError(
                f"offset >= 0 and 4 <= limit <= {albedo_api.OUTPUT_PREVIEW};"
                " page with offset="
            )
        data = self._bytes(offset, limit)
        skip = 0
        while (
            skip < min(3, len(data))
            and offset + skip < self.size
            and data[skip] & 0xC0 == 0x80
        ):
            skip += 1  # a page that opens inside a character starts at the next
        data = data[skip:]
        if offset + skip + len(data) < self.size:
            data = data[: albedo_output.whole_characters(data)]
        return albedo_api.Text(data.decode("utf-8", errors="replace"))

    def head(self, n: int = 4000, *, lines: int | None = None) -> Any:
        self._require_text()
        data = self._bytes(0, albedo_api.OUTPUT_PREVIEW)
        return albedo_api.excerpt(
            data.decode("utf-8", errors="ignore"),
            n,
            lines,
            end=False,
            clipped=self.size > len(data),
        )

    def tail(self, n: int = 4000, *, lines: int | None = None) -> Any:
        self._require_text()
        start = max(0, self.size - albedo_api.OUTPUT_PREVIEW)
        data = self._bytes(start, albedo_api.OUTPUT_PREVIEW)
        return albedo_api.excerpt(
            data.decode("utf-8", errors="ignore"),
            n,
            lines,
            end=True,
            clipped=start > 0,
        )

    def __iter__(self) -> Iterator[Any]:
        self._require_text()
        if self.kind == "bytes":
            raise TypeError("bytes on disk have no lines; page them with .read()")
        with open(self.path, encoding="utf-8", newline="") as file:
            for line in file:
                yield line.rstrip("\n") if self.kind == "lines" else line

    def load(self) -> Any:
        """The value, whole. Refused when it is past the memory cap."""
        if self.size > albedo_memory.LIMIT:
            raise MemoryError(
                f"{self.path} is {mebibytes(self.size)}, past the kernel's"
                f" {mebibytes(albedo_memory.LIMIT)} memory cap; page it with .read(),"
                " iterate it, or files.read(path, start_line=, end_line=)"
            )
        if self.kind == "pickle":
            with open(self.path, "rb") as file:
                return cast("Any", albedo_state.engine()[0]).load(file)
        with open(self.path, "rb") as file:
            data = file.read()
        if self.kind == "bytes":
            return data
        text = data.decode("utf-8")
        if self.kind == "text":
            return text
        lines = text.split("\n")[:-1] if text else []
        return tuple(lines) if self.container == "tuple" else lines


def _destination(name: str, suffix: str) -> tuple[BinaryIO, str]:
    handle, path = tempfile.mkstemp(
        prefix=f"value-{name}-", suffix=suffix, dir=albedo_capture.SPILL_DIR
    )
    return os.fdopen(handle, "wb"), path


def _write_text(file: BinaryIO, text: str) -> None:
    for start in range(0, len(text), WRITE_CHUNK):
        _ = file.write(text[start : start + WRITE_CHUNK].encode("utf-8"))


def _is_lines(value: object) -> bool:
    return (
        isinstance(value, (list, tuple))
        and bool(value)
        and all(
            isinstance(item, str) and "\n" not in item and "\r" not in item
            for item in value
        )
    )


def move(name: str, value: object) -> Paged | None:
    """`value` written to the output directory, or None when it cannot be
    (an unpicklable object, or no directory)."""
    if albedo_capture.SPILL_DIR is None:
        return None
    length = len(value) if isinstance(value, Sized) else None
    if isinstance(value, str):
        kind, suffix, container = "text", ".txt", ""
    elif isinstance(value, (bytes, bytearray)):
        kind, suffix, container = "bytes", ".bin", ""
    elif _is_lines(value):
        kind, suffix, container = "lines", ".txt", type(value).__name__
    else:
        kind, suffix, container = "pickle", ".pkl", type(value).__name__
    try:
        file, path = _destination(name, suffix)
    except OSError:
        return None
    try:
        with file:
            if isinstance(value, str):
                _write_text(file, value)
            elif isinstance(value, (bytes, bytearray)):
                _ = file.write(value)
            elif kind == "lines":
                lines = cast("list[str]", value)
                for start in range(0, len(lines), 100_000):
                    _write_text(file, "\n".join(lines[start : start + 100_000]) + "\n")
            else:
                cast("Any", albedo_state.engine()[0]).dump(value, file)
        return Paged(path, kind, os.path.getsize(path), length, container)
    except BaseException:
        try:
            os.unlink(path)
        except OSError:
            pass
        return None


def trim() -> None:
    """Gives freed memory back to the system where the allocator can."""
    _ = gc.collect()
    try:
        ctypes.CDLL("libc.so.6").malloc_trim(0)
    except (OSError, AttributeError):
        pass


def relieve(namespace: dict[str, object], injected: set[str]) -> str:
    """When the kernel is past its cap, moves every large value the namespace
    binds to disk and returns the note that names them, or "" when nothing
    moved. The moved names are rebound to `Paged` handles."""
    used = albedo_memory.footprint()
    if used is None or used <= albedo_memory.LIMIT:
        return ""
    held = sorted(
        (
            (size_of(value), name)
            for name, value in list(namespace.items())
            if not name.startswith("_")
            and name not in injected
            and not isinstance(
                value, (types.ModuleType, types.FunctionType, type, Paged)
            )
        ),
        key=lambda item: item[0],
        reverse=True,
    )
    moved: list[str] = []
    handles: dict[int, Paged] = {}
    value: object = None
    with albedo_memory.headroom():
        for size, name in held:
            if size < MOVE_MIN:
                break
            value = namespace[name]
            paged = handles.get(id(value)) or move(name, value)
            if paged is None:
                continue
            handles[id(value)] = paged
            namespace[name] = paged
            moved.append(f"{name} ({mebibytes(size)})")
        del value
        handles.clear()
    if not moved:
        return ""
    trim()
    now = albedo_memory.footprint() or 0
    stays = (
        " The process stays above the cap until its session rests and it is recycled."
        if now > albedo_memory.LIMIT
        else ""
    )
    return (
        f"[kernel memory was {mebibytes(used)}, cap {mebibytes(albedo_memory.LIMIT)}:"
        f" {', '.join(moved)} moved to disk, {mebibytes(now)} held now. Page a value"
        " with name.read(offset, limit), name.head() or name.tail(lines=), iterate it"
        f" for its lines, or name.load() (refused past the cap).{stays}]\n"
    )