"""Utilities to smuggle data into stdout.

The basic idea is to print out a (syntactically valid but meaningless) escape
sequence that terminals will ignore, that way the output looks fine on the
terminal, but there is extra data embedded in the stream that can be recovered.
"""

import re
from collections import deque

# Standard hexadecimal characters
oldhex = "0123456789abcdef"

# Arbitrary sequences of these sixteen characters are valid in escape sequences
newhex = "0123456789:;<=>?"

hexsub = dict(zip(oldhex, newhex))
hexsub_back = dict(zip(newhex, oldhex))


def encode_as_escape_sequence(s: str):
    """Encode a string as a CSI escape sequence.

    The escape sequence is meaningless, but terminals will ignore it, so we can
    smuggle data invisibly that way.

    A CSI code starts with the ESC character followed by [, an arbitrary number
    of characters in the 0x30 to 0x3F range (see the newhex variable), then 0 or more
    characters in the 0x20 to 0x2F range and then one byte from 0x40-0x7E.

    We simply encode the data as hex using the characters in 0x30-0x3F, to which we
    append the letter z to end the escape sequence.

    Arguments:
        s: The string to encode.
    """
    encoded = "".join([hexsub[c] for c in s.encode("utf8").hex()])
    return f"\033[{encoded}z"


def decode_escape_sequence(s):
    """Decode an escape sequence generated by :func:`encode_as_escape_sequence`."""
    s = s[2:-1]
    decoded_hex = "".join([hexsub_back[c] for c in s])
    return bytes.fromhex(decoded_hex).decode("utf8")


class SmuggleWriter:
    """Class for smuggling data into a stream using escape sequences.

    This uses :func:`encode_as_escape_sequence` on each write.

    Arguments:
        stream: The stream in which to smuggle the data, e.g. ``sys.stdout``
    """

    def __init__(self, stream):
        self.stream = stream

    def __enter__(self):
        self.stream.__enter__()
        return self

    def __exit__(self, typ=None, exc=None, tb=None):
        self.stream.__exit__(typ, exc, tb)
        return self

    def write(self, txt):
        return self.stream.write(encode_as_escape_sequence(txt))

    def flush(self):
        return self.stream.flush()

    def close(self):
        return self.stream.close()


class LineAccumulator:
    """Process individual characters and accumulate full lines into a deque."""

    def __init__(self):
        self.lines = deque()
        self.current = ""

    def process(self, char):
        self.current += char
        if char == "\n":
            self.lines.append(self.current)
            self.current = ""


def utf8_length(char: chr) -> int:
    """Return number of bytes a UTF-8 character uses from its first byte."""
    b = ord(char)

    if b & 0b10000000 == 0:
        return 1  # ASCII
    elif b & 0b11100000 == 0b11000000:
        return 2
    elif b & 0b11110000 == 0b11100000:
        return 3
    elif b & 0b11111000 == 0b11110000:
        return 4
    return 1


class Decoder:
    """Decoder for a stream where data has been smuggled.

    Produces two :class:`LineAccumulator`s, ``out`` and ``data``:

    * ``out`` contains the normal output of the stream.
    * ``data`` contains everything that was encoded through
      :func:`encode_as_escape_sequence`

    Arguments:
        principal: The stream.
    """

    def __init__(self, principal):
        self.principal = principal
        self.out = LineAccumulator()
        self.data = LineAccumulator()
        self.code = ""
        self.coding = False

    def close(self):
        self.principal.close()

    def endcode(self):
        self.coding = False
        if re.match(string=self.code, pattern="\033\\[[0-9:;<=>?]*z"):
            for char in decode_escape_sequence(self.code):
                self.data.process(char)
        else:
            self.out.current += self.code
        self.code = ""

    def process_char(self, char):
        if self.coding:
            if ord(char) < 0x20:
                self.out.process(char)
                self.endcode()
            elif ord(char) >= 0x40 and char != "[":
                self.code += char
                self.endcode()
            else:
                self.code += char

        elif char == "\033":
            self.coding = True
            self.code += char

        else:
            self.out.process(char)

    def getline(self, which):
        if which == "out":
            if self.out.lines:
                return self.out.lines.popleft()
        elif which == "data":
            if self.data.lines:
                return self.data.lines.popleft()
        return None

    def readline(self, which: str):
        """Get the next line from either ``self.out`` or ``self.data``.

        Arguments:
            which: Either ``"out"`` or ``"data"``.

        Returns:
            The next line or ``None`` if there is no data.
        """
        while (result := self.getline(which)) is None:
            nxt = self.principal.read(1)
            if not nxt:
                return None

            for i in range(utf8_length(nxt) - 1):
                nxt += self.principal.read(1)

            self.process_char(nxt.decode("utf8"))

        return result


class MultimodalFile:
    """File object for a stream in which data has been smuggled.

    Arguments:
        decoder: The :class:`Decoder` instance to read from.
        which: Either ``"out"`` or ``"data"``.
    """

    def __init__(self, decoder: Decoder, which: str, name: str):
        self.name = name
        self.decoder = decoder
        self.which = which

    def close(self):
        self.decoder.close()

    def fileno(self):
        return self.name

    def readline(self):
        return self.decoder.readline(self.which)
