"""Low-level PDF objects and operations"""
from __future__ import annotations

import abc
import re
from binascii import hexlify
from dataclasses import dataclass
from functools import partial
from itertools import accumulate, chain, repeat, starmap
from math import isfinite
from secrets import token_bytes
from typing import Collection, Iterable, Iterator, Sequence, Tuple
from zlib import compress

from .common import add_slots, setattr_frozen

# PDF 32000-1:2008 (7.5.2) specifies that the header must be followed
# immediately by a comment line containing at least 4 binary characters
# of size 0x80 or greater.
_HEADER = b"%PDF-1.7\n%\x80\x80\x80\x80 generated by pdfje\n"
_FIRST_OFFSET = len(_HEADER)
OBJ_ID_XREF, OBJ_ID_CATALOG, OBJ_ID_PAGETREE, OBJ_ID_RESOURCES = 0, 1, 2, 3
ASCII = bytes
Byte = int  # 0-255
# ID of an object in the PDF file. Always >=0 and unique within a file.
ObjectID = int


class Atom(abc.ABC):
    __slots__ = ()

    @abc.abstractmethod
    def write(self) -> Iterable[bytes]:
        raise NotImplementedError()


Object = Tuple[ObjectID, Atom]


@add_slots
@dataclass(frozen=True)
class Bool(Atom):
    value: bool

    def write(self) -> Iterable[bytes]:
        raise NotImplementedError()


@add_slots
@dataclass(frozen=True)
class Name(Atom):
    value: ASCII

    def write(self) -> Iterable[bytes]:
        return (b"/", self.value)


def _sanitize_name_char(c: Byte) -> bytes:
    # PDF32000-1:2008 (7.3.5) says ASCII from 33-126 is OK, except "#" (0x23)
    # which is the escape character.
    if c == 0x23:
        return b"#23"
    elif 33 <= c <= 126:
        return c.to_bytes(1, "big")
    # We decide to keep spaces, but remove the rest
    elif c == 0x20:
        return b"#20"
    else:
        return b""


def sanitize_name(s: bytes) -> bytes:
    return b"".join(map(_sanitize_name_char, s))


@add_slots
@dataclass(frozen=True)
class Int(Atom):
    value: int

    def write(self) -> Iterable[bytes]:
        return (b"%i" % self.value,)


@add_slots
@dataclass(frozen=True)
class Real(Atom):
    value: float

    def write(self) -> Iterable[bytes]:
        assert isfinite(self.value), "NaN and Inf not supported"
        return (b"%g" % self.value,)


# NOTE: this name avoids confusion with typing.LiteralString
@add_slots
@dataclass(frozen=True)
class LiteralStr(Atom):
    "See PDF32000-1:2008 (7.3.4.2)"
    value: bytes

    # FUTURE: support UTF-16BOM, but in a way that makes it explicit that
    # this does not support usage in text streams.
    # How to prevent accidential UTF-16BOM?
    def write(self) -> Iterable[bytes]:
        return (b"(", _escape(self.value), b")")


@add_slots
@dataclass(frozen=True)
class HexString(Atom):
    value: bytes

    def write(self) -> Iterable[bytes]:
        return (b"<", hexlify(self.value), b">")


@add_slots
@dataclass(frozen=True)
class Array(Atom):
    items: Iterable[Atom]

    def write(self) -> Iterable[bytes]:
        yield b"["
        for i in self.items:
            yield from i.write()
            yield b" "
        yield b"]"


@add_slots
@dataclass(frozen=True, init=False)
class Dictionary(Atom):
    content: Collection[tuple[ASCII, Atom]]

    def __init__(self, *content: tuple[ASCII, Atom]) -> None:
        setattr_frozen(self, "content", content)

    def write(self) -> Iterable[bytes]:
        yield from _write_dict(self.content)


def _write_dict(d: Iterable[tuple[ASCII, Atom]]) -> Iterable[bytes]:
    yield b"<<\n"
    for key, value in d:
        yield b"/"
        yield key
        yield b" "
        yield from value.write()
        yield b"\n"
    yield b">>"


@add_slots
@dataclass(frozen=True)
class Stream(Atom):
    content: Iterable[bytes]
    meta: Collection[tuple[ASCII, Atom]] = ()

    def write(self) -> Iterable[bytes]:
        content = compress(b"".join(self.content))
        yield from _write_dict(
            chain(
                self.meta,
                [
                    (b"Length", Int(len(content))),
                    (b"Filter", Name(b"FlateDecode")),
                ],
            )
        )
        yield b"\nstream\n"
        yield content
        yield b"\nendstream"


@add_slots
@dataclass(frozen=True)
class Ref(Atom):
    target: ObjectID

    def write(self) -> Iterable[bytes]:
        return (b"%i 0 R" % self.target,)


def _write_obj(i: ObjectID, o: Atom) -> Iterable[bytes]:
    yield b"%i 0 obj\n" % i
    yield from o.write()
    yield b"\nendobj\n"


# PDF32000:1-2008 (14.4) says the file ID must be unique byte string
# identifying the file. They recommend MD5 hashing the contents,
# file location, and time.
# We'll take a shortcut here and just use random, which will do fine.
# We use `secrets` to ensure this randomness cannot be manipulated.
_generate_file_id = partial(token_bytes, 16)


def _write_trailer(
    # offsets must be already sorted continuously ascending by object ID from 1
    offsets: Sequence[int],
    xref_offset: int,
) -> Iterable[bytes]:
    yield b"xref\n%i %i\n0000000000 65535 f \n" % (
        OBJ_ID_XREF,
        len(offsets) + 1,
    )
    yield from map(b"%010i 00000 n \n".__mod__, offsets)
    yield b"trailer\n"
    yield from _write_dict(
        (
            (b"Root", Ref(OBJ_ID_CATALOG)),
            (b"Size", Int(len(offsets) + 1)),
            (b"ID", Array(repeat(HexString(_generate_file_id()), 2))),
        )
    )
    yield b"\nstartxref\n%i\n%%%%EOF\n" % xref_offset


def write(
    # object IDs must ascend continuously from 4 onwards,
    # with the last three items being catalog, pagetree,
    # and resources (ids 1, 2, 3 respectively).
    objs: Iterable[Object],
) -> Iterator[bytes]:
    yield _HEADER
    offsets = [_FIRST_OFFSET]
    for chunk in map(b"".join, starmap(_write_obj, objs)):
        offsets.append(len(chunk))
        yield chunk
    (
        *offsets,
        catalog_offset,
        pagetree_offset,
        resources_offset,
        xref_offset,
    ) = accumulate(offsets)
    yield from _write_trailer(
        [catalog_offset, pagetree_offset, resources_offset] + offsets,
        xref_offset,
    )


_STRING_ESCAPES = {
    b"\\": b"\\\\",
    b"\n": b"\\n",
    b"\r": b"\\r",
    b"\t": b"\\t",
    b"\b": b"\\b",
    b"\f": b"\\f",
    b"(": b"\\(",
    b")": b"\\)",
}


def _replace_with_escape(m: re.Match[bytes]) -> bytes:
    return _STRING_ESCAPES[m.group()]


_escape = partial(
    re.compile(b"(%b)" % b"|".join(map(re.escape, _STRING_ESCAPES))).sub,
    _replace_with_escape,
)
