#!/usr/bin/python3

import asyncio
import aiohttp
import platform
import logging
import json
import os
import shlex
import subprocess
from multiprocessing import cpu_count
from psutil import virtual_memory, disk_usage
from netifaces import ifaddresses, AF_INET
from launchy import Launchy

# allow older aiohttp versions
if not hasattr(aiohttp, "client_exceptions") and hasattr(aiohttp, "errors"):
    aiohttp.client_exceptions = aiohttp.errors


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("molior-client")
molior_server = os.environ.get("MOLIOR_SERVER", "172.16.0.254")
interface_name = os.environ.get("INTERFACE_NAME", "eth0")


async def build(params, masterws):
    ret = -1
    try:
        build_id = params.get("build_id")
        apt_urls = params.get("apt_urls")
        apt_keys = params.get("apt_keys")
        token = params.get("token")

        logger.info("starting build_%d", build_id)

        await masterws.send_str(json.dumps({"status": "building", "build_id": build_id}))

        sbuild_apt_urls = ["--extra-repository=\"{}\"".format(url) for url in apt_urls]
        # set env for build script
        env = os.environ.copy()
        env["BUILD_ID"] = str(build_id)
        env["BUILD_TOKEN"] = token
        env["REPO_NAME"] = params.get("repository_name")
        env["VERSION"] = params.get("version")
        env["PLATFORM"] = params.get("distrelease")
        env["PLATFORM_VERSION"] = params.get("distversion")
        env["ARCH"] = params.get("architecture")
        env["ARCH_ANY_ONLY"] = str(int(params.get("arch_any_only")))
        env["APT_URLS"] = " ".join(sbuild_apt_urls)
        env["APT_KEYS"] = " ".join(apt_keys)
        env["APT_URLS_PLAIN"] = "\n".join([" - {}".format(url) for url in apt_urls])
        env["APT_URLS_PLAIN"] += "\n"
        env["PROJECT"] = params.get("project")
        env["PROJECTVERSION"] = params.get("projectversion")
        env["APT_SERVER"] = params.get("apt_server")
        env["PROJECT_DIST"] = params.get("project_dist")
        env["RUN_LINTIAN"] = "1" if params.get("run_lintian", False) else "0"

        session = aiohttp.ClientSession()
        async with session.ws_connect("ws://{}/internal/buildlog/{}".format(molior_server, token)) as buildws:

            async def wslogger():
                try:
                    async for msg in buildws:
                        if msg.type == aiohttp.WSMsgType.TEXT:
                            logger.info("got message: %s", msg.data)
                        elif msg.type == aiohttp.WSMsgType.CLOSED:
                            break
                        elif msg.type == aiohttp.WSMsgType.ERROR:
                            break
                except Exception as exc:
                    logger.exception(exc)

            async def output(data):
                await buildws.send_str(data)

            buildcmd = "/usr/bin/unbuffer /usr/lib/molior/build-script"

            try:
                process = Launchy(shlex.split(buildcmd), output, output, buffered=False, collect_time=0.1, env=env)
                await process.launch()
                ret = await process.wait()
            except Exception as exc:
                logger.exception(exc)

        logger.info("build-script returned %d", ret)
    except Exception as exc:
        logger.error("Error running build script")
        logger.exception(exc)

    try:
        await masterws.send_str(json.dumps({"status": "success" if ret == 0 else "failed", "build_id": build_id}))
    except Exception as exc:
        logger.error("Error sending build status")
        logger.exception(exc)
        masterws.close()


async def main():

    def get_machine_id():
        try:
            with open('/etc/machine-id') as machine_id_file:
                machine_id = machine_id_file.readline().strip()
        except IOError:
            machine_id = None
        return machine_id

    def get_ip_address():
        interface = {}
        inet_config = ifaddresses(interface_name).get(AF_INET, None)
        if inet_config:
            if len(inet_config) > 0:
                interface = inet_config[0]
        ip = interface.get("addr", "")
        if ip == "?":
            return None
        return ip

    session = aiohttp.ClientSession()
    machine = platform.machine()
    node = platform.node()
    cpu_cores = cpu_count()
    ram_total = virtual_memory().total
    disk_total = disk_usage("/").total
    machine_id = get_machine_id()
    client_ver = str(subprocess.check_output(["dpkg-query", "--showformat=${Version}", "--show",
                     "molior-client-http"], stderr=subprocess.DEVNULL), "utf-8")

    if machine == 'x86_64':
        arch = 'amd64'
    elif machine == 'aarch64':
        arch = 'arm64'
    else:
        logger.error("invalid machine architecture: '%s'", machine)
        return

    logger.info("starting on %s/%s", arch, node)

    while(True):
        await asyncio.sleep(1)
        try:
            async with session.ws_connect("ws://{}/internal/registry/{}/{}".format(
                                          molior_server, arch, node)) as ws:
                await ws.send_str(json.dumps({"register": {"cpu_cores": cpu_cores,
                                                           "ram_total": ram_total,
                                                           "disk_total": disk_total,
                                                           "id": machine_id,
                                                           "ip": get_ip_address(),
                                                           "client_ver": client_ver}}))
                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        try:
                            req = json.loads(msg.data)
                        except Exception as exc:
                            logger.error("invalid request: %s", msg.data)
                            logger.exception(exc)
                            await asyncio.sleep(1)
                            continue

                        if "task" in req:
                            asyncio.create_task(build(req["task"], ws))
                        elif "ping" in req:
                            uptime_seconds = ""
                            with open('/proc/uptime', 'r') as f:
                                uptime_seconds = float(f.readline().split()[0])

                            await ws.send_str(json.dumps({"pong": {
                                              "seq": req["ping"],
                                              "uptime_seconds": uptime_seconds,
                                              "load": os.getloadavg(), "ram_used": virtual_memory().used,
                                              "disk_used": disk_usage("/").used}}))

                        else:
                            logger.error("invalid request: %s", msg.data)
                    elif msg.type == aiohttp.WSMsgType.CLOSED:
                        break
                    elif msg.type == aiohttp.WSMsgType.ERROR:
                        break
        except aiohttp.client_exceptions.WSServerHandshakeError as exc:
            logger.info("WebSocket error: %s", str(exc.message))
            await asyncio.sleep(4)
        except aiohttp.client_exceptions.ClientConnectorError:
            await asyncio.sleep(4)
        except Exception as exc:
            logger.exception(exc)

if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    Launchy.attach_loop(loop)
    asyncio.ensure_future(main())
    loop.run_forever()
