import asyncio
import asyncssh
from asgiref.sync import async_to_sync
from channels.generic.websocket import AsyncJsonWebsocketConsumer
from channels.consumer import AsyncConsumer
from channels.layers import get_channel_layer


class WebShellSocketConsumer(AsyncJsonWebsocketConsumer):
    def __init__(self):
        super().__init__()

        self.conn = None

    async def connect(self):
        super().connect()

        if self.scope['session'].session_key is None:
            await self.scope['session'].acreate()
            await self.scope['session'].asave()

        await self.accept()

        print('accepted connection as', self.scope['session'].session_key)

    async def receive_json(self, data):
        data['owner'] = self.scope['session'].session_key

        await self.channel_layer.group_add(data['owner'], self.channel_name)
        await self.channel_layer.send('terminal', data)

    async def notify(self, event):
        await self.send_json({
            'type': 'stdout',
            'stdout': event['data']
        })


class WebShellClient(asyncssh.SSHClient):
    pass


class WebShellClientSession(asyncssh.SSHClientSession):
    def data_received(self, data, datatype):
        channel = get_channel_layer()

        asyncio.ensure_future(channel.group_send(self.channel_name, {
            'type': 'notify',
            'data': data,
        }))

    def eof_received(self):
        print('ssh connection done')


class WebShellWorker(AsyncConsumer):
    def __init__(self):
        self.sessions = {}

    async def ssh(self, event):
        host = event['host']
        user = event['username']
        pw = event['password']

        if not host:
            await self.send_json({
                'error': 'no IP provided'
            })

            return

        if not user:
            await self.send_json({
                'error': 'no username provided'
            })

            return

        if not pw:
            await self.send_json({
                'error': 'no password provided'
            })

            return

        # TODO
        # Track these on Consumer

        if not event['owner'] in self.sessions:
            print('creating new session')

            conn, client = await asyncssh.create_connection(
                WebShellClient,
                host,
                username=user,
                password=pw,
                known_hosts=None,
                options=asyncssh.SSHClientConnectionOptions(
                    #server_host_key_algs='ssh-rsa'
                )
            )

            chan, session = await conn.create_session(
                WebShellClientSession,
                env={
                    'channel_name': event['owner'],
                },
                term_type='xterm'
            )

            session.channel_name = event['owner']
            session.worker = self

            self.sessions[event['owner']] = {
                'connection': conn,
                'client': client,
                'channel': chan,
                'session': session,
            }

    async def key(self, event):
        self.sessions[event['owner']]['channel'].write(event['key'])