# -*- coding: utf-8 -*-
import os
import json
import random
import time
import base64
import threading
from tornado import ioloop
from tornado import gen
from tornado.websocket import websocket_connect
from tornado.tcpclient import TCPClient

class ClientSide:

    def __init__(self, local_ip, remote_ip):
        self._local_ip = local_ip
        self._remote_ip = remote_ip
        self.local_conn = None
        self.local_read_future = None
        self.remote_conn = None
        self.conn_hash = {}

    @gen.coroutine
    def co_localConnect(self, idx):
        print(idx, flush=True)
        [local_host, local_port] = self._local_ip.split(':')
        local_conn = yield TCPClient().connect(local_host, local_port)
        self.conn_hash[idx] = local_conn
        self.sendRemoteMsg({
            'cmd' : 'connectOk',
            'remoteConnIdx' : idx,
        })
        while True:
            try:
                data = yield local_conn.read_until_close()
                if data:
                    self.sendRemoteMsg({
                        'cmd' : 'forwardData',
                        'remoteConnIdx' : idx,
                        'data' : str(base64.b64encode(data))
                    })
                break
            except:
                print('qqqqq', flush=True)
                break

    @gen.coroutine
    def co_remoteConnect(self):
        [remote_host, remote_port] = self._remote_ip.split(':')
        url = 'ws://%s:%s/websocket' % (remote_host, remote_port)
        self.remote_conn = yield websocket_connect(url)
        data = ''
        try:
            while True:
                data += yield self.remote_conn.read_message()
                if not data:
                    continue
                lines = data.split('\n')
                if data[-1] == '\n':
                    data = lines[-1]
                    lines = lines[:-1]
                for line in lines:
                    msg = json.loads(line)
                    self.dispatchRemoteMsg(msg)
        except:
            print('wwwwwwwwww', flush=True)


    def dispatchRemoteMsg(self, msg):
        if msg['cmd'] == 'connect':
            ioloop.IOLoop.current().spawn_callback(self.co_localConnect, msg['remoteConnIdx'])
        elif msg['cmd'] == 'socketClose':
            pass
        elif msg['cmd'] == 'forwardData':
            if msg['remoteConnIdx'] in self.conn_hash:
                conn = self.conn_hash[msg['remoteConnIdx']]
                data = base64.b64decode(msg['data'][2:-1])
                print(data, flush=True)
                conn.write(data)

    def sendRemoteMsg(self, msg):
        data = json.dumps(msg) + '\n'
        self.remote_conn.write_message(data)

    @gen.coroutine
    def co_connect(self):
#        ioloop.IOLoop.current().spawn_callback(self.co_localConnect)
        ioloop.IOLoop.current().spawn_callback(self.co_remoteConnect)
        while True:
            yield gen.sleep(0.1)

    def run(self):
        ioloop.IOLoop.current().run_sync(self.co_connect)