#!/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)

'''
    {
      "md5": "6c7257dd9ee9342f376f4b140df1f9a0",
      "version": "4.3.28.11361",
      "size": 7250217,
      "path": "U7P/4.3.28.11361/a1e0-U7P-4.3.28-62a146a0d8dd4161a291a74fd562b9ce.bin",
      "devices": [
        "U7P"
      ]
    }
'''

'''
      {
        "channel": "release",
        "created": "2021-02-10T09:03:12Z",
        "file_size": 7252057,
        "id": "3a44dd11-6c86-45ef-8d4e-5dbb91349ea7",
        "md5": "02da6330a4cf868d02819bce4d583651",
        "sha256_checksum": "255736db010c226897c10bb32818992401aaf9b942910caafae7c08d5fab06a7",
        "platform": "BZ2",
        "product": "unifi-firmware",
        "updated": "2021-02-10T09:03:15Z",
        "version": "v4.3.28+11361",
        "version_major": 4,
        "version_minor": 3,
        "version_patch": 28,
        "version_build": "11361",
        "_links": {
          "self": {
            "href": "https://fw-update.ubnt.com/api/firmware/3a44dd11-6c86-45ef-8d4e-5dbb91349ea7"
          },
          "upload": [
            {
              "name": "data",
              "href": "https://fw-update.ubnt.com/api/firmware/3a44dd11-6c86-45ef-8d4e-5dbb91349ea7/data"
            },
            {
              "name": "changelog",
              "href": "https://fw-update.ubnt.com/api/firmware/3a44dd11-6c86-45ef-8d4e-5dbb91349ea7/changelog"
            }
          ],
          "data": {
            "href": "https://fw-download.ubnt.com/data/unifi-firmware/8822-BZ2-4.3.28-3a44dd116c8645ef8d4e5dbb91349ea7.bin"
          }
        }
      },
'''