import datetime
import logging
import pickle
import queue
import socket
import threading
import traceback
from multiprocessing import shared_memory

import numpy as np

from .serializer import dumper
from .serializer import loader
from instamatic import config
from instamatic.camera import Camera
from instamatic.utils import high_precision_timers
high_precision_timers.enable()


condition = threading.Condition()
box = []

HOST = config.settings.cam_server_host
PORT = config.settings.cam_server_port
BUFSIZE = 4096


is_local_connection = HOST in ('127.0.0.1', 'localhost')


class CamServer(threading.Thread):
    """Camera communcation server.

    Takes a logger object `log`, command queue `q`, and name of the
    camera `name` that is used to initialize the connection to the
    camera. Start the server using `CamServer.run` which will wait for
    items to appear on `q` and execute them on the specified camera
    instance.
    """

    def __init__(self, log=None, q=None, name=None):
        super().__init__()

        self.log = log
        self.q = q

        # self.name is a reserved parameter for threads
        self._name = name

        self._bufsize = BUFSIZE

        self.verbose = False

        self.buffers = {}

        self.use_shared_memory = config.settings.cam_use_shared_memory
        print('Use shared memory:', self.use_shared_memory)

    def setup_shared_buffer(self, arr):
        """Set up shared memory buffer.

        Make a buffer for each binsize, and store the buffers to a dict.
        """
        self.shmem = shared_memory.SharedMemory(create=True, size=arr.nbytes)
        buffer = np.ndarray(arr.shape, dtype=arr.dtype, buffer=self.shmem.buf)
        self.buffers[arr.shape] = buffer
        if self.verbose:
            print(f'Created new buffer: `{self.shmem.name}` | {arr.shape} ({arr.dtype})')

    def copy_data_to_shared_buffer(self, arr):
        """Copy numpy image array to shared memory."""
        if arr.shape not in self.buffers:
            self.setup_shared_buffer(arr)

        buffer = self.buffers[arr.shape]
        buffer[:] = arr[:]  # copy data to buffer

    def run(self):
        """Start server thread."""
        self.cam = Camera(name=self._name, use_server=False)
        self.cam.get_attrs = self.get_attrs

        print(f'Initialized camera: {self.cam.interface}')

        while True:
            now = datetime.datetime.now().strftime('%H:%M:%S.%f')

            cmd = self.q.get()

            with condition:
                attr_name = cmd['attr_name']
                args = cmd.get('args', ())
                kwargs = cmd.get('kwargs', {})

                try:
                    ret = self.evaluate(attr_name, args, kwargs)
                    status = 200
                except Exception as e:
                    traceback.print_exc()
                    if self.log:
                        self.log.exception(e)
                    ret = (e.__class__.__name__, e.args)
                    status = 500
                else:
                    if self.use_shared_memory:
                        if attr_name == 'getImage':
                            self.copy_data_to_shared_buffer(ret)
                            ret = {
                                'shape': ret.shape,
                                'dtype': str(ret.dtype),
                                'name': self.shmem.name,
                            }

                box.append((status, ret))
                condition.notify()
                if self.verbose:
                    print(f'{now} | {status} {attr_name}: {ret}')

    def evaluate(self, attr_name: str, args: list, kwargs: dict):
        """Evaluate the function or attribute `attr_name` on `self.cam`, if
        `attr_name` refers to a function, call it with *args and **kwargs."""
        # print(attr_name, args, kwargs)
        f = getattr(self.cam, attr_name)
        if callable(f):
            ret = f(*args, **kwargs)
        else:
            ret = f
        return ret

    def get_attrs(self):
        """Get attributes from cam object to update __dict__ on client side."""
        attrs = {}
        for item in dir(self.cam):
            if item.startswith('_'):
                continue
            obj = getattr(self.cam, item)
            if not callable(obj):
                attrs[item] = type(obj)

        return attrs


def handle(conn, q):
    """Handle incoming connection, put command on the Queue `q`, which is then
    handled by TEMServer."""
    with conn:
        while True:
            data = conn.recv(BUFSIZE)
            if not data:
                break

            data = loader(data)

            if data == 'exit':
                break

            if data == 'kill':
                break

            with condition:
                q.put(data)
                condition.wait()
                response = box.pop()
                conn.sendall(dumper(response))


def main():
    import argparse

    description = f"""
Connects to the camera and starts a server for camera communication. Opens a socket on port {HOST}:{PORT}.

This program initializes a connection to the camera as defined in the config. This separates the communication from the main process and allows for remote connections from different PCs. The connection goes over a TCP socket.

The host and port are defined in `config/settings.yaml`.

The data sent over the socket is a pickled dictionary with the following elements:

- `attr_name`: Name of the function to call or attribute to return (str)
- `args`: (Optional) List of arguments for the function (list)
- `kwargs`: (Optiona) Dictionary of keyword arguments for the function (dict)

The response is returned as a pickle object.
"""

    parser = argparse.ArgumentParser(
        description=description,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument('-c', '--camera', action='store', dest='camera',
                        help="""Override camera to use.""")

    parser.set_defaults(camera=None)
    options = parser.parse_args()
    camera = options.camera

    date = datetime.datetime.now().strftime('%Y-%m-%d')
    logfile = config.logs_drc / f'instamatic_CAMServer_{date}.log'
    logging.basicConfig(format='%(asctime)s | %(module)s:%(lineno)s | %(levelname)s | %(message)s',
                        filename=logfile,
                        level=logging.DEBUG)
    logging.captureWarnings(True)
    log = logging.getLogger(__name__)

    q = queue.Queue(maxsize=100)

    cam_reader = CamServer(name=camera, log=log, q=q)
    cam_reader.start()

    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind((HOST, PORT))
    s.listen(5)

    log.info(f'Server listening on {HOST}:{PORT}')
    print(f'Server listening on {HOST}:{PORT}')

    with s:
        while True:
            conn, addr = s.accept()
            log.info('Connected by %s', addr)
            print('Connected by', addr)
            threading.Thread(target=handle, args=(conn, q)).start()


if __name__ == '__main__':
    main()
