#!/usr/bin/python2
############################################################
import argparse
import sys
import os
import zipfile
import json
import apt_inst
import onlu
import subprocess

logger = onlu.init_logging("switool")


class OnlSwitchImage(object):
    def __init__(self, fname, mode):
        self.fname = fname
        self.mode = mode
        # use allowZip64 to support rootfs over 2GB
        self.zipfile = zipfile.ZipFile(fname, mode=mode, allowZip64=True)
        self.manifest = None

    def add(self, fname, arcname=None, compressed=True):
        self.zipfile.write(
            fname,
            arcname=arcname,
            compress_type=zipfile.ZIP_DEFLATED if compressed else zipefile.ZIP_STORED,
        )

    def add_rootfs(self, rootfs_sqsh):
        self.add(rootfs_sqsh)

    def add_manifest(self, manifest):
        self.add(manifest, arcname="manifest.json")

    def get_manifest(self):
        if "manifest.json" in self.zipfile.namelist():
            return json.load(self.zipfile.open("manifest.json"))
        else:
            return None

    def get_arch(self):
        return self.get_manifest()["arch"]

    def get_platforms(self):
        p = self.get_manifest()["platforms"]
        if type(p) is list:
            return p
        else:
            return p.split(",")

    def get_contents(self):
        return self.zipfile.namelist()


############################################################

ap = argparse.ArgumentParser(description="SWI Tool")
ap.add_argument("--create", action="store_true", help="Create new SWI.")
ap.add_argument("--overwrite", action="store_true", help="Overwrite existing file.")
ap.add_argument("--rootfs", help="Root SquashFS File")
ap.add_argument("--manifest", help="SWI Manifest file.")
ap.add_argument("--add-files", help="Add additional files.", default=[], nargs="+")
ap.add_argument("--contents", help="Show SWI contents.", action="store_true")
ap.add_argument("--platforms", help="Show SWI contents.", action="store_true")
ap.add_argument("swi", help="SWI image name.")

ops = ap.parse_args()

if os.path.exists(ops.swi):
    if ops.create and not ops.overwrite:
        logger.critical("File '%s' exists." % ops.swi)
        sys.exit(1)

swi = None

if ops.create or ops.overwrite:
    if not ops.rootfs:
        logger.critical("Rootfs required to create new SWI.")
        sys.exit(1)

    if not ops.manifest:
        logger.critical("Manifest required to create new SWI.")
        sys.exit(1)

    swi = OnlSwitchImage(ops.swi, "w")
    swi.add_rootfs(ops.rootfs)
    swi.add_manifest(ops.manifest)
    for f in ops.add_files:
        swi.add(f, arcname=f)

if swi is None:

    if not os.path.exists(ops.swi):
        logger.critical("SWI file %s does not exist." % ops.swi)
        sys.exit(1)

    swi = OnlSwitchImage(ops.swi, "r")

if ops.contents:
    print(" ".join(swi.get_contents()))

if ops.platforms:
    print(" ".join(swi.get_platforms()))