#!/usr/bin/env python3

import _jsonnet as j
import json
import yaml
import logging
import os
import sys
import zipfile
import pathlib
from io import BytesIO

logger = logging.getLogger("generate")
logging.basicConfig(level=logging.INFO, format='%(message)s')

private_json = "Put your GCP private.json here"

class Generator:

    def __init__(
        self, config, templates="./templates/", resources="./resources",
        version="0.0.0",
    ):

        self.templates = pathlib.Path(templates)
        self.resources = pathlib.Path(resources)
        self.config = config
        self.version = f"\"{version}\"".encode("utf-8")

    def process(self, config):

        res = j.evaluate_snippet("config", config, import_callback=self.load)
        return json.loads(res)

    def load(self, dir, filename):

        logger.debug("Request jsonnet: %s %s", dir, filename)

        if filename == "config.json" and dir == "":
            path = os.path.join(".", dir, filename)
            return str(path), self.config

        if filename == "version.jsonnet" and dir == "templates/values/":
            path = os.path.join(".", dir, filename)
            return str(path), self.version

        if dir:
            candidates = [
                self.templates.joinpath(dir, filename),
                self.templates.joinpath(filename),
                self.resources.joinpath(dir, filename),
                self.resources.joinpath(filename),
                pathlib.Path(dir).joinpath(filename),
            ]
        else:
            candidates = [
                self.templates.joinpath(filename),
                pathlib.Path(dir).joinpath(filename),
                pathlib.Path(filename),
            ]

        try:

            if filename == "vertexai/private.json":

                return str(candidates[0]), private_json.encode("utf-8")

            for c in candidates:
                logger.debug("Try: %s", c)

                if os.path.isfile(c):
                    with open(c, "rb") as f:
                        logger.debug("Loading: %s", c)
                        return str(c), f.read()

            raise RuntimeError(
                f"Could not load file={filename} dir={dir}"
            )
                
        except:

            path = os.path.join(self.templates, filename)
            logger.debug("Try: %s", path)
            with open(path, "rb") as f:
                logger.debug("Loaded: %s", path)
                return str(path), f.read()

class Packager:

    def __init__(self):
        self.templates = pathlib.Path("./templates")
        self.resources = pathlib.Path("./")

    def process(
            self, config, version="0.0.0", platform="docker-compose",
    ):

        config = config.encode("utf-8")

        gen = Generator(
            config, templates=self.templates, resources=self.resources,
            version=version
        )

        path = self.templates.joinpath(
            f"config-to-{platform}.jsonnet"
        )
        wrapper = path.read_text()

        processed = gen.process(wrapper)

        return processed
    
    def generate(self, config, version, platform):

        logger.info(f"Generating for platform={platform} version={version}")

        try:

            if platform in set(["docker-compose", "podman-compose"]):
                return self.generate_docker_compose(
                    "docker-compose", version, config
                )
            elif platform in set(["minikube-k8s", "gcp-k8s"]):
                return self.generate_k8s(
                    platform, version, config
                )
            else:
                raise RuntimeError("Bad configuration")

        except Exception as e:
            logging.error(f"Exception: {e}")
            raise e

    def generate_docker_compose(self, platform, version, config):

        processed = self.process(
            config, platform=platform, version=version
        )

        y = yaml.dump(processed)

        mem = BytesIO()

        with zipfile.ZipFile(mem, mode='w') as out:

            def output(name, content):
                logger.info(f"Adding {name}...")
                out.writestr(name, content)

            fname = "docker-compose.yaml"

            output(fname, y)

            # Grafana config
            path = self.resources.joinpath(
                "grafana/dashboards/dashboard.json"
            )
            res = path.read_text()
            output("grafana/dashboards/dashboard.json", res)

            path = self.resources.joinpath(
                "grafana/provisioning/dashboard.yml"
            )
            res = path.read_text()
            output("grafana/provisioning/dashboard.yml", res)

            path = self.resources.joinpath(
                "grafana/provisioning/datasource.yml"
            )
            res = path.read_text()
            output("grafana/provisioning/datasource.yml", res)

            # Prometheus config
            path = self.resources.joinpath(
                "prometheus/prometheus.yml"
            )
            res = path.read_text()
            output("prometheus/prometheus.yml", res)

        logger.info("Generation complete.")

        return mem.getvalue()

    def generate_k8s(self, platform, version, config):

        processed = self.process(
            config, platform=platform, version=version
        )

        y = yaml.dump(processed)

        mem = BytesIO()

        with zipfile.ZipFile(mem, mode='w') as out:

            def output(name, content):
                logger.info(f"Adding {name}...")
                out.writestr(name, content)

            fname = "resources.yaml"

            output(fname, y)

        logger.info("Generation complete.")

        return mem.getvalue()
            
def main():

    if len(sys.argv) != 4:
        print()
        print("Usage:")
        print("  generate <outfile> <version> <platform> < input.json")
        print()
        sys.exit(1)

    outfile = sys.argv[1]
    version = sys.argv[2]
    platform = sys.argv[3]

    cfg = sys.stdin.read()
    
    logger.info(f"Outputting to {outfile}...")

    p = Packager()
    resp = p.generate(cfg, version, platform)

    with open(outfile, "wb") as f:
        f.write(resp)

    return

main()