import base64
import contextlib
import hashlib
import io
import os
import logging
import queue
import select
import socket
import struct
import subprocess
import sys
import threading
import time

from enum import Enum
from http.server import BaseHTTPRequestHandler
from itertools import cycle

from .proto import messages_pb2 as pb2

logger = logging.getLogger(__name__)

class NAL:
    CODED_SLICE_NON_IDR = 1  # Coded slice of a non-IDR picture
    CODED_SLICE_IDR     = 5  # Coded slice of an IDR picture
    SEI                 = 6  # Supplemental enhancement information (SEI)
    SPS                 = 7  # Sequence parameter set
    PPS                 = 8  # Picture parameter set

ALLOWED_NALS = {NAL.CODED_SLICE_NON_IDR,
                NAL.CODED_SLICE_IDR,
                NAL.SPS,
                NAL.PPS,
                NAL.SEI}

def StartMessage(resolution):
    width, height = resolution
    return pb2.ClientBound(timestamp_us=int(time.monotonic() * 1000000),
                           start=pb2.Start(width=width, height=height))

def StopMessage():
    return pb2.ClientBound(timestamp_us=int(time.monotonic() * 1000000),
                           stop=pb2.Stop())

def VideoMessage(data):
    return pb2.ClientBound(timestamp_us=int(time.monotonic() * 1000000),
                           video=pb2.Video(data=data))

def OverlayMessage(svg):
    return pb2.ClientBound(timestamp_us=int(time.monotonic() * 1000000),
                           overlay=pb2.Overlay(svg=svg))

def _parse_server_message(data):
    message = pb2.ServerBound()
    message.ParseFromString(data)
    return message

def _shutdown(sock):
    try:
        sock.shutdown(socket.SHUT_RDWR)
    except OSError:
        pass

def _read_asset(path):
    if path == '/':
        path = 'index.html'
    elif path[0] == '/':
        path = path[1:]

    base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'assets'))
    asset_path = os.path.abspath(os.path.join(base_path, path))

    if os.path.commonpath((base_path, asset_path)) != base_path:
        return None, None

    if path.endswith('.html'):
        content_type = 'text/html; charset=utf-8'
    elif path.endswith('.js'):
        content_type = 'application/javascript; charset=utf-8'
    elif path.endswith('.wasm'):
        content_type = 'application/wasm'
    else:
        content_type = 'application/octet-stream'

    try:
        with open(asset_path, 'rb') as f:
            return f.read(), content_type
    except Exception:
        return None, None


class HTTPRequest(BaseHTTPRequestHandler):

    def __init__(self, request_buf):
        self.rfile = io.BytesIO(request_buf)
        self.raw_requestline = self.rfile.readline()
        self.parse_request()


def _read_http_request(sock):
    request = bytearray()
    while b'\r\n\r\n' not in request:
        buf = sock.recv(2048)
        if not buf:
            break
        request.extend(buf)
    return request


def _http_ok(content, content_type):
    header = (
        'HTTP/1.1 200 OK\r\n'
        'Content-Length: %d\r\n'
        'Content-Type: %s\r\n'
        'Connection: Keep-Alive\r\n\r\n'
    ) % (len(content), content_type)
    return header.encode('ascii') + content


def _http_switching_protocols(token):
    accept_token = token.encode('ascii') + b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
    accept_token = hashlib.sha1(accept_token).digest()
    header = (
        'HTTP/1.1 101 Switching Protocols\r\n'
        'Upgrade: WebSocket\r\n'
        'Connection: Upgrade\r\n'
        'Sec-WebSocket-Accept: %s\r\n\r\n'
    ) % base64.b64encode(accept_token).decode('ascii')
    return header.encode('ascii')


def _http_not_found():
    return 'HTTP/1.1 404 Not Found\r\n\r\n'.encode('ascii')


@contextlib.contextmanager
def Socket(port):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(('', port))
    sock.listen()
    try:
        yield sock
    finally:
        _shutdown(sock)
        sock.close()


class DroppingQueue:

    def __init__(self, maxsize):
        if maxsize <= 0:
            raise ValueError('Maxsize must be positive.')
        self.maxsize = maxsize
        self._items = []
        self._cond = threading.Condition(threading.Lock())

    def put(self, item, replace_last=False):
        with self._cond:
            was_empty = len(self._items) == 0
            if len(self._items) < self.maxsize:
                self._items.append(item)
                if was_empty:
                    self._cond.notify()
                return False  # Not dropped.

            if replace_last:
                self._items[len(self._items) - 1] = item
                return False  # Not dropped.

            return True  # Dropped.

    def get(self):
        with self._cond:
            while not self._items:
                self._cond.wait()
            return self._items.pop(0)


class AtomicSet:

    def __init__(self):
        self._lock = threading.Lock()
        self._set = set()

    def add(self, value):
        with self._lock:
            self._set.add(value)
            return value

    def remove(self, value):
        with self._lock:
            try:
                self._set.remove(value)
                return True
            except KeyError:
                return False

    def __len__(self):
        with self._lock:
            return len(self._set)

    def __iter__(self):
        with self._lock:
            return iter(self._set.copy())

class PresenceServer:

    SERVICE_TYPE = '_aiy_vision_video._tcp'

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def __init__(self, name, port):
        logger.info('Start publishing %s on port %d.', name, port)
        cmd = ['avahi-publish-service', name, self.SERVICE_TYPE, str(port), 'AIY Streaming']
        self._process = subprocess.Popen(cmd, shell=False)

    def close(self):
        self._process.terminate()
        self._process.wait()
        logger.info('Stop publishing.')


class StreamingServer:

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def __init__(self, camera, bitrate=1000000, mdns_name=None,
                 tcp_port=4665, web_port=4664, annexb_port=4666):
        self._bitrate = bitrate
        self._camera = camera
        self._clients = AtomicSet()
        self._enabled_clients = AtomicSet()
        self._done = threading.Event()
        self._commands = queue.Queue()
        self._thread = threading.Thread(target=self._run,
                                        args=(mdns_name, tcp_port, web_port, annexb_port))
        self._thread.start()

    def close(self):
        self._done.set()
        self._thread.join()

    def send_overlay(self, svg):
        for client in self._enabled_clients:
            client.send_overlay(svg)

    def _start_recording(self):
        logger.info('Camera start recording')
        self._camera.start_recording(self, format='h264', profile='baseline',
            inline_headers=True, bitrate=self._bitrate, intra_period=0)

    def _stop_recording(self):
        logger.info('Camera stop recording')
        self._camera.stop_recording()

    def _process_command(self, client, command):
        was_streaming = bool(self._enabled_clients)

        if command is ClientCommand.ENABLE:
            self._enabled_clients.add(client)
        elif command is ClientCommand.DISABLE:
            self._enabled_clients.remove(client)
        elif command == ClientCommand.STOP:
            self._enabled_clients.remove(client)
            if self._clients.remove(client):
                client.stop()
            logger.info('Number of active clients: %d', len(self._clients))

        is_streaming = bool(self._enabled_clients)

        if not was_streaming and is_streaming:
            self._start_recording()
        if was_streaming and not is_streaming:
            self._stop_recording()

    def _run(self, mdns_name, tcp_port, web_port, annexb_port):
        try:
            with contextlib.ExitStack() as stack:
                logger.info('Listening on ports tcp: %d, web: %d, annexb: %d',
                            tcp_port, web_port, annexb_port)
                tcp_socket = stack.enter_context(Socket(tcp_port))
                web_socket = stack.enter_context(Socket(web_port))
                annexb_socket = stack.enter_context(Socket(annexb_port))
                if mdns_name:
                    stack.enter_context(PresenceServer(mdns_name, tcp_port))

                socks = (tcp_socket, web_socket, annexb_socket)
                while not self._done.is_set():
                    # Process available client commands.
                    try:
                        while True:
                            client, command = self._commands.get_nowait()
                            self._process_command(client, command)
                    except queue.Empty:
                        pass  # Done processing commands.

                    # Process recently connected clients.
                    rlist, _, _ = select.select(socks, [], [], 0.2)  # 200ms
                    for ready in rlist:
                        sock, addr = ready.accept()
                        name = '%s:%d' % addr
                        if ready is tcp_socket:
                            client = ProtoClient(name, sock, self._commands, self._camera.resolution)
                        elif ready is web_socket:
                            client = WsProtoClient(name, sock, self._commands, self._camera.resolution)
                        elif ready is annexb_socket:
                            client = AnnexbClient(name, sock, self._commands)
                        logger.info('New %s connection from %s', client.TYPE, name)

                        self._clients.add(client).start()
                        logger.info('Number of active clients: %d', len(self._clients))
        finally:
            logger.info('Server is shutting down')
            if self._enabled_clients:
                self._stop_recording()

            for client in self._clients:
                client.stop()
            logger.info('Done')

    def write(self, data):
        """Called by camera thread for each compressed frame."""
        assert data[0:4] == b'\x00\x00\x00\x01'
        frame_type = data[4] & 0b00011111
        if frame_type in ALLOWED_NALS:
            states = {client.send_video(frame_type, data) for client in self._enabled_clients}
            if ClientState.ENABLED_NEEDS_SPS in states:
                logger.info('Requesting key frame')
                self._camera.request_key_frame()

class ClientLogger(logging.LoggerAdapter):
    def process(self, msg, kwargs):
        return '[%s] %s' % (self.extra['name'], msg), kwargs

class ClientState(Enum):
    DISABLED = 1
    ENABLED_NEEDS_SPS = 2
    ENABLED = 3

class ClientCommand(Enum):
    STOP = 1
    ENABLE = 2
    DISABLE = 3

class Client:
    def __init__(self, name, sock, command_queue):
        self._lock = threading.Lock()  # Protects _state.
        self._state = ClientState.DISABLED
        self._logger = ClientLogger(logger, {'name': name})
        self._socket = sock
        self._commands = command_queue
        self._tx_q = DroppingQueue(15)
        self._rx_thread = threading.Thread(target=self._rx_run)
        self._tx_thread = threading.Thread(target=self._tx_run)

    def start(self):
        self._rx_thread.start()
        self._tx_thread.start()

    def stop(self):
        self._logger.info('Stopping...')
        _shutdown(self._socket)
        self._socket.close()
        self._tx_q.put(None)
        self._tx_thread.join()
        self._rx_thread.join()
        self._logger.info('Stopped.')

    def send_video(self, frame_type, data):
        """Only called by camera thread."""
        with self._lock:
            if self._state == ClientState.DISABLED:
                pass
            elif self._state == ClientState.ENABLED_NEEDS_SPS:
                if frame_type == NAL.SPS:
                    dropped = self._queue_video(data)
                    if not dropped:
                        self._state = ClientState.ENABLED
            elif self._state == ClientState.ENABLED:
                dropped = self._queue_video(data)
                if dropped:
                    self._state = ClientState.ENABLED_NEEDS_SPS
            return self._state

    def send_overlay(self, svg):
        """Can be called by any user thread."""
        with self._lock:
            if self._state != ClientState.DISABLED:
                self._queue_overlay(svg)

    def _send_command(self, command):
        self._commands.put((self, command))

    def _queue_message(self, message, replace_last=False):
        dropped = self._tx_q.put(message, replace_last)
        if dropped:
            self._logger.warning('Running behind, dropping messages')
        return dropped

    def _tx_run(self):
        try:
            while True:
                message = self._tx_q.get()
                if message is None:
                    break
                self._send_message(message)
            self._logger.info('Tx thread finished')
        except Exception as e:
            self._logger.warning('Tx thread failed: %s', e)

        # Tx thread stops the client in any situation.
        self._send_command(ClientCommand.STOP)

    def _rx_run(self):
        try:
            while True:
                message = self._receive_message()
                if message is None:
                    break
                self._handle_message(message)
            self._logger.info('Rx thread finished')
        except Exception as e:
            self._logger.warning('Rx thread failed: %s', e)
            # Rx thread stops the client only if error happened.
            self._send_command(ClientCommand.STOP)

    def _receive_bytes(self, num_bytes):
        received = bytearray()
        while len(received) < num_bytes:
            buf = self._socket.recv(num_bytes - len(received))
            if not buf:
                return buf
            received.extend(buf)
        return received

    def _queue_video(self, data):
        raise NotImplementedError

    def _queue_overlay(self, svg):
        raise NotImplementedError

    def _send_message(self, message):
        raise NotImplementedError

    def _receive_message(self):
        raise NotImplementedError

    def _handle_message(self, message):
        pass

class ProtoClient(Client):
    TYPE = 'tcp'

    def __init__(self, name, sock, command_queue, resolution):
        super().__init__(name, sock, command_queue)
        self._resolution = resolution

    def _queue_video(self, data):
        return self._queue_message(VideoMessage(data))

    def _queue_overlay(self, svg):
        return self._queue_message(OverlayMessage(svg))

    def _handle_message(self, message):
        which = message.WhichOneof('message')
        if which == 'stream_control':
            self._handle_stream_control(message.stream_control)

    def _handle_stream_control(self, stream_control):
        enabled = stream_control.enabled
        self._logger.info('stream_control %s', enabled)

        with self._lock:
            if self._state == ClientState.DISABLED and not enabled:
                self._logger.info('Ignoring stream_control disable')
            elif self._state in (ClientState.ENABLED_NEEDS_SPS, ClientState.ENABLED) and enabled:
                self._logger.info('Ignoring stream_control enable')
            else:
                if enabled:
                    self._logger.info('Enabling client')
                    self._state = ClientState.ENABLED_NEEDS_SPS
                    self._queue_message(StartMessage(self._resolution))
                    self._send_command(ClientCommand.ENABLE)
                else:
                    self._logger.info('Disabling client')
                    self._state = ClientState.DISABLED
                    self._queue_message(StopMessage(), replace_last=True)
                    self._send_command(ClientCommand.DISABLE)

    def _send_message(self, message):
        buf = message.SerializeToString()
        self._socket.sendall(struct.pack('!I', len(buf)))
        self._socket.sendall(buf)

    def _receive_message(self):
        buf = self._receive_bytes(4)
        if not buf:
            return None
        num_bytes = struct.unpack('!I', buf)[0]
        buf = self._receive_bytes(num_bytes)
        if not buf:
            return None
        return _parse_server_message(buf)


class WsProtoClient(ProtoClient):
    TYPE = 'web'

    class WsPacket:
        def __init__(self):
            self.fin = True
            self.opcode = 2
            self.masked = False
            self.mask = None
            self.length = 0
            self.payload = bytearray()

        def append(self, data):
            if self.masked:
                data = bytes([c ^ k for c, k in zip(data, cycle(self.mask))])
            self.payload.extend(data)

        def serialize(self):
            self.length = len(self.payload)
            buf = bytearray()
            b0 = 0
            b1 = 0
            if self.fin:
                b0 |= 0x80
            b0 |= self.opcode
            buf.append(b0)
            if self.length <= 125:
                b1 |= self.length
                buf.append(b1)
            elif self.length >= 126 and self.length <= 65535:
                b1 |= 126
                buf.append(b1)
                buf.extend(struct.pack('!H', self.length))
            else:
                b1 |= 127
                buf.append(b1)
                buf.extend(struct.pack('!Q', self.length))
            if self.payload:
                buf.extend(self.payload)
            return bytes(buf)

    def __init__(self, name, sock, command_queue, resolution):
        super().__init__(name, sock, command_queue, resolution)
        self._upgraded = False

    def _receive_message(self):
        try:
            if not self._upgraded:
                if self._process_web_request():
                    return None
                self._upgraded = True

            packets = []
            while True:
                packet = self._receive_packet()
                if packet.opcode == 0:
                    # Continuation
                    if not packets:
                        self._logger.error('Invalid continuation received')
                        return None
                    packets.append(packet)
                elif packet.opcode == 1:
                    # Text, not supported.
                    self._logger.error('Received text packet')
                    return None
                elif packet.opcode == 2:
                    # Binary.
                    packets.append(packet)
                    if packet.fin:
                        joined = bytearray()
                        for p in packets:
                            joined.extend(p.payload)
                        return _parse_server_message(joined)
                elif packet.opcode == 8:
                    # Close.
                    self._logger.info('WebSocket close requested')
                    return None
                elif packet.opcode == 9:
                    # Ping, send pong.
                    self._logger.info('Received ping')
                    response = self.WsPacket()
                    response.opcode = 10
                    response.append(packet.payload)
                    self._queue_message(response)
                elif packet.opcode == 10:
                    # Pong. Igore as we don't send pings.
                    self._logger.info('Dropping pong')
                else:
                    self._logger.info('Dropping opcode %d', packet.opcode)
        except Exception:
            self._logger.exception('Error while processing websocket request')
            return None

    def _receive_packet(self):
        packet = self.WsPacket()
        buf = self._receive_bytes(2)
        packet.fin = buf[0] & 0x80 > 0
        packet.opcode = buf[0] & 0x0F
        packet.masked = buf[1] & 0x80 > 0
        packet.length = buf[1] & 0x7F
        if packet.length == 126:
            packet.length = struct.unpack('!H', self._receive_bytes(2))[0]
        elif packet.length == 127:
            packet.length = struct.unpack('!Q', self._receive_bytes(8))[0]
        if packet.masked:
            packet.mask = self._receive_bytes(4)
        packet.append(self._receive_bytes(packet.length))
        return packet

    def _send_message(self, message):
        if isinstance(message, (bytes, bytearray)):
            buf = message
        else:
            if isinstance(message, self.WsPacket):
                packet = message
            else:
                packet = self.WsPacket()
                packet.append(message.SerializeToString())
            buf = packet.serialize()
        self._socket.sendall(buf)

    def _process_web_request(self):
        request = _read_http_request(self._socket)
        request = HTTPRequest(request)
        connection = request.headers['Connection']
        upgrade = request.headers['Upgrade']
        if 'Upgrade' in connection and upgrade == 'websocket':
            sec_websocket_key = request.headers['Sec-WebSocket-Key']
            self._queue_message(_http_switching_protocols(sec_websocket_key))
            self._logger.info('Upgraded to WebSocket')
            return False

        if request.command == 'GET':
            content, content_type = _read_asset(request.path)
            if content is None:
                self._queue_message(_http_not_found())
            else:
                self._queue_message(_http_ok(content, content_type))
            self._queue_message(None)
            return True

        raise Exception('Unsupported request')


class AnnexbClient(Client):
    TYPE = 'annexb'

    def __init__(self, name, sock, command_queue):
        super().__init__(name, sock, command_queue)
        self._state = ClientState.ENABLED_NEEDS_SPS
        self._send_command(ClientCommand.ENABLE)

    def _queue_video(self, data):
        return self._queue_message(data)

    def _queue_overlay(self, svg):
        pass  # Ignore overlays.

    def _send_message(self, message):
        self._socket.sendall(message)

    def _receive_message(self):
        buf = self._socket.recv(1024)
        if not buf:
            return None
        raise RuntimeError('Invalid state.')