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 .helpers.old_wifi import search_wifi, disconnect, connect_wifi
from .helpers.read_and_write import write_json, delete_json, json_parse
from .helpers import  check_outdated
from enum import Enum
from http.server import BaseHTTPRequestHandler
from itertools import cycle
import json
from .proto import messages_pb2 as pb2
import Orange_Vision
from urllib import request as urllib
# is_outdated, latest_version = check_outdated('Orange_Vision', Orange_Vision.__version__)
# print(is_outdated, latest_version)


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 USBMessage(data):
    return pb2.ClientBound(timestamp_us=int(time.monotonic() * 1000000),
                           usb=pb2.Usb(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 _file_content_type(path):
    if path.endswith('.html'):
        return 'text/html; charset=utf-8'
    elif path.endswith('.js'):
        return 'text/javascript; charset=utf-8'
    elif path.endswith('.css'):
        return 'text/css; charset=utf-8'
    elif path.endswith('.png'):
        return'image/png'
    elif path.endswith('.svg'):
        return 'image/svg+xml; charset=utf-8'
    elif path.endswith('.jpg') or path.endswith('.jpeg'):
        return'image/jpeg'
    elif path.endswith('.wasm'):
        return'application/wasm'
    else:
        return 'application/octet-stream'

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

def _asset_path(path):
    
    if path == '/':
        value = os.environ.get('SERVER_INDEX_HTML')
        if value is not None:
            return value
        path  = 'index.html'
    elif path == '/wifi':
        logger.info("wifi connect")
        path  = 'test.html'
    elif path == '/settings':
        logger.info("settings")
        path  = 'settings.html'
    elif path[0] == '/':
        path = path[1:]

    asset_path = os.path.abspath(os.path.join(BASE_PATH, path))
    if os.path.commonpath((BASE_PATH, asset_path)) != BASE_PATH:
        return None

    return asset_path

def _read_asset(path):
    asset_path = _asset_path(path)
    if asset_path is not None:
        with contextlib.suppress(Exception):
            with open(asset_path, 'rb') as f:
                return f.read(), _file_content_type(asset_path)
    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=5801, 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()
                
    def write_usb(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_usb_vid(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_usb_vid(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_usb_video(data)
                    if not dropped:
                        self._state = ClientState.ENABLED
            elif self._state == ClientState.ENABLED:
                dropped = self._queue_usb_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_usb_video(self, data):
        return self._queue_message(USBMessage(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 send_post_data(self, data):
        content = json.dumps(data).encode('utf-8')

        content_type = 'text/plain; charset=utf-8'
        self._queue_message(_http_ok(content, content_type))
        self._queue_message(None)

    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 == 'POST':


            if request.path == '/getVersion':
                self.send_post_data(Orange_Vision.__version__.strip('\"'))
                
            if request.path == '/updateAndReboot':
                os.system("sudo apt-get update && sudo apt-get dist-upgrade -y")
                os.system("pip3 install --upgrade Orange_Vision")
                os.system("sudo reboot now")
            if request.path == '/checkUpdate':

                try:
                    url = "https://www.google.com"
                    urllib.urlopen(url)
                    status = "Connected"
                except:
                    status = "Not connected"
                print(status, "status")
                if status == "Connected":
                    is_outdated, latest_version = check_outdated('Orange_Vision', Orange_Vision.__version__)
                    data = {"outdated": is_outdated, "latest_version": latest_version}
                    self.send_post_data(data)
                
                else:
                    data = "False"
                    self.send_post_data(data)


                
            if request.path == '/disconnect':
                path = request.headers['Authority']
                print(path, "disconnect")
                disconnect(path)
                data = "False"
                self.send_post_data(data)

            if request.path == '/connect':
                path = request.headers['Authority']
                data = json.loads(path)
                result = connect_wifi(data)
                self.send_post_data(result)


            if request.path == '/getConnections':
                wifi_list = search_wifi()
                
                self.send_post_data(wifi_list)
            

            if request.path == '/add_profile':
                data = request.headers['Authority']
                data = json.loads(data)
                write_json('profiles', data)


            if request.path == '/delete_profile':
                data = request.headers['Authority']
                data = json.loads(data)
                data = delete_json('profiles', data)

            return True



        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.')
