__author__ = 'JackSong'

import logging

import tornado.websocket

from daemon import Bridge
from utils.data import ClientData
from utils.tools import check_ip, check_port
from concurrent.futures import ThreadPoolExecutor
from tornado.gen import coroutine
from utils.tools import PrpcryptInitData


class IndexHandler(tornado.web.RequestHandler):
    def get(self):
        self.render("index.html")


class WSHandler(tornado.websocket.WebSocketHandler):
    executor = ThreadPoolExecutor(4)
    clients = dict()
    COMMAND = str()

    def get_client(self):
        return self.clients.get(self._id(), None)

    def put_client(self):
        bridge = Bridge(self)
        self.clients[self._id()] = bridge

    def remove_client(self):
        bridge = self.get_client()
        if bridge:
            client = self.clients.pop(self._id(), None)
            if client: del client

    @staticmethod
    def _check_init_param(data):
        data = PrpcryptInitData(**data)
        return check_ip(data.hostname) and check_port(data.port)

    @staticmethod
    def _is_init_data(data):
        return data.get_type() == 'init'

    def _id(self):
        return id(self)

    def open(self):
        self.put_client()

    @coroutine
    def on_message(self, message):
        bridge = self.get_client()
        client_data = ClientData(message)
        if self._is_init_data(client_data):
            if self._check_init_param(client_data.data):
                try:
                    yield self.executor.submit(bridge.open, *(client_data.data,))
                    logging.info('connection established from: %s' % self._id())
                except Exception as e:
                    self.close()
                    logging.error('%s' % e.message)

            else:
                self.remove_client()
                logging.info('init param invalid: %s' % client_data.data)
        else:
            if bridge.status:
                command = client_data.data
                rest = yield self.executor.submit(bridge.trans_forward, *(command,))
                if not rest:
                    bridge.destroy()
                    self.close()

    def on_close(self):
        self.remove_client()
        logging.info('client close the connection: %s' % self._id())

    def check_origin(self, origin):
        return True