#!/usr/bin/env python3

"""
    This is a low level script that communicate with a device over the network.

    Currently only a swupdate-www-push subcommand is available.

    It can be used this way:
        coreos-device swupdate-www-push ~/coreos-image-all-features-vm-x64.swu 192.168.122.21

    This script doesn't interact with bitbake so it can't retrieve variables from
    bitbake. For a more user friendly usage you can use devtool:

    devtool swupdate-www-push coreos-image-all-features 192.168.122.21
"""

import os
import sys
import pathlib
import ipaddress
import asyncio
import json

# We need some modules that are not part of the Python Standard Library
# If they are not found, let's fail with an appropriate error message
try:
    import click
    import aiohttp
except ModuleNotFoundError as e:
    print(f"Error: {e}", file=sys.stderr)
    print(f"Please install the python3-{e.name} package", file=sys.stderr)
    sys.exit(1)

@click.group()
def cli():
    pass

@cli.command()
@click.argument( 'image', type=click.Path(
    exists=True, resolve_path=True, readable=True, dir_okay=False,
    path_type=pathlib.Path,
))
@click.argument('device', type=ipaddress.ip_address)
@click.option('--port', type=click.IntRange(min=1, max=65535), default=8080)
def swupdate_www_push(image, device, port):
    loop = asyncio.new_event_loop()
    returncode = loop.run_until_complete(swupdate_www_push_async(image, device, port))
    loop.close()

    if returncode == 0:
        click.secho("coreos-device: Update was successfull", bg="green")
    else:
        click.secho(f"coreos-device: Update has failed!", bg="red")

    sys.exit(returncode)


async def swupdate_www_push_async(image, device, port):
    url = f'http://{device}:{port}'

    async with aiohttp.ClientSession() as session:
        tasks = []
        
        tasks.append(asyncio.ensure_future(swupdate_www_push_data(session, url, image)))
        tasks.append(asyncio.ensure_future(swupdate_www_websocket_logger(session, url)))

        responses = await asyncio.gather(*tasks)

        # Return the sum of return value of all ours tasks
        returncode = 0
        for response in responses:
            returncode += response
        return returncode
        

async def swupdate_www_push_data(session, url, image):

    with open(image, 'rb') as f:
        data = {'file': f}
        try:
            async with session.post(f"{url}/upload", data=data) as resp:
                if resp.status == 200:
                    click.secho(f"coreos-device: http status {resp.status}", bg="white", fg="black")
                    return 0
                else:
                    click.secho(f"coreos-device: http status {resp.status}, {await resp.text()}", bg="red")
                    return resp.status
        except Exception as e:
            click.secho(f"coreos-device: ERROR: {e}", bg="red")
            return 1        


async def swupdate_www_websocket_logger(session, url):
    return_code = 0
    try:
        async with session.ws_connect(f"{url}/ws") as ws:
            async for msg in ws:
                data = json.loads(msg.data)
                if data["type"] == "message":
                    msg = data["text"]
                    if "Waiting for requests" in msg:
                        await ws.close()
                        return 1
                    elif "ERROR" in msg or "not successful" in msg:
                        click.secho("swupdate: " + data["text"], fg='red')
                        return_code = 1
                    elif "successful" in msg:
                        click.secho("swupdate: " + data["text"], fg="green")
                    else:
                        click.secho("swupdate: " + data["text"])
                    
                if data["type"] == "status":
                    if data["status"] == "DONE":
                        await ws.close()
                        return return_code
    except Exception as e:
        click.secho(f"coreos-device: ERROR: {e}", bg="red")
        return 1 

if __name__ == '__main__':
    cli()