#!/usr/bin/env python

import logging
import os
import json
import requests

def load_json(filename):
    logging.info(f"reading json file {filename}")
    try:
        with open(filename) as f:
            data = json.load(f)
            return data
    except:
        logging.exception(f"could not read file {filename}")
    return None

def write_json(filename, data):
    logging.info(f"writing json file {filename}")
    try:
        with open(filename, "w") as f:
            json.dump(data, f)
    except:
        logging.exception(f"could not write file {filename}")

def load_firmware_meta(path):
    try:
        firmware_meta = load_json(os.path.join(path, "firmware_meta.json"))
        if firmware_meta is None:
            return []
        return firmware_meta.get('cached_firmwares', [])
    except:
        logging.exception(f"could load firmware meta")
        return []
    return firmware_meta

def write_firmware_meta(path, data):
    try:
        firmware_meta = {
            'cached_firmwares': data
        }
        write_json(os.path.join(path, "firmware_meta.json"), firmware_meta)
    except:
        logging.exception(f"could write firmware meta")

def online_firmware_check():
    try:
        resp = requests.get("https://fw-update.ubnt.com/api/firmware-newest",
            params={"filter": ["eq~~product~~unifi-firmware", "eq~~channel~~release"]})
        data = resp.json()
        return data.get("_embedded").get("firmware", [])
    except:
        logging.exception(f"could load firmware updates")
        return []

def online_firmware_download(url, filename):
    try:
        path = os.path.dirname(filename)
        os.makedirs(path, exist_ok=True)
        response = requests.get(url, stream=True)
        with open(filename, mode="wb") as file:
            for chunk in response.iter_content(chunk_size=10 * 1024):
                file.write(chunk)
    except:
        logging.exception(f"could not download firmware {url}")
        return False
    else:
        return True

def main(args):
    firmware_meta = load_firmware_meta(args.firmware_dir)
    firmware_latest = online_firmware_check()

    # loop over the all available firmware
    # if found in 'meta', just add device to ther list
    # we assume, the firmware is already downloaded
    # if not found in 'meta', add ne item and download firmware

    for firmware in firmware_latest:
        firmware_platform = firmware.get('platform')
        firmware_version = [
            str(firmware.get('version_major')),
            str(firmware.get('version_minor')),
            str(firmware.get('version_patch')),
            str(firmware.get('version_build')),
        ]
        firmware_version = ".".join(firmware_version)

        if args.platform is not None:
            if firmware_platform not in args.platform:
                continue

        logging.info(f"processing firmware {firmware_platform}/{firmware_version}")

        found = None

        for meta in firmware_meta:
            # compare some values
            if meta.get('version') != firmware_version:
                continue
            if meta.get('md5') != firmware.get('md5'):
                continue
            if meta.get('size') != firmware.get('file_size'):
                continue
            # if all values match, we got our item
            found = meta
            break

        # just add device to the list
        if found is not None:
            if firmware.get('platform') not in found.get('devices'):
                found['devices'].append(firmware.get('platform'))
            continue

        firmware_url = firmware.get('_links', {}).get('data', {}).get('href', {})
        if firmware_url is None:
            logging.error(f"unalbe to get firmware path for {firmware_platform}/{firmware_version}")
            continue
        found = {
            'md5': firmware.get('md5'),
            'version': firmware_version,
            'size': firmware.get('file_size'),
            'path': os.path.join(firmware.get('platform'), firmware_version,firmware_url.split("/")[-1]),
            'devices': [firmware.get('platform')]
        }
        if online_firmware_download(firmware_url, os.path.join(args.firmware_dir, found['path'])):
            firmware_meta.append(found)

    write_firmware_meta(args.firmware_dir, firmware_meta)

if __name__ == "__main__":

    format = "%(asctime)s: %(message)s"
    logging.basicConfig(format=format, level=logging.DEBUG, datefmt="%H:%M:%S")

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--firmware-dir", type=str, required=True, help="directory with firmwares")
    parser.add_argument("-p","--platform", action='append', help='platform list to download', required=False)
    args = parser.parse_args()
    main(args)