#!/usr/bin/env python

from __future__ import print_function

import os, sys, tempfile, subprocess, shutil, re, argparse
import yaml

parser = argparse.ArgumentParser(description='Pack a robot.zip from user code.')
parser.add_argument('team_tla',
                    nargs='?',
                    help='The TLA of the team the zip is being created for')
parser.add_argument('usercode',
                    help='user code source directory, containing a robot.py')
parser.add_argument('output',
                    help='output zip file')
parser.add_argument("-w", "--wifi",
                    default="/etc/sr/wifi-keys.yaml",
                    help="YAML file to get wifi keys from")

args = parser.parse_args()

USER_CODE_DIR, OUTPUT_ZIP = args.usercode, args.output

if not os.path.isdir( USER_CODE_DIR ):
    print("Error: Supplied user directory path is not a directory.",
          file=sys.stderr)
    exit(1)

REPO_ROOT = os.path.dirname( __file__ )
BASE_DIR = os.path.join( REPO_ROOT, "base" )
assert os.path.isdir( BASE_DIR )

# Temporary directory for constructing everything inside
tmpd = tempfile.mkdtemp( suffix="-pyenv" )

subprocess.check_call( "cp -r %s/* %s" % (BASE_DIR, tmpd),
                       shell = True )

# Copy in the user's code
subprocess.check_call( ["cp", "-r", USER_CODE_DIR, os.path.join( tmpd, "user" )] )

# Do the wifi key dance
default_wifi = True
if os.path.exists(args.wifi) and args.team_tla is not None:
    with open(args.wifi, "r") as f:
        y = yaml.safe_load(f)

    if args.team_tla in y:
        with open(os.path.join(tmpd, "wifi.yaml"), "w") as f:
            yaml.safe_dump( { "ssid": "robot-{}".format(args.team_tla),
                              "key": y[args.team_tla] }, f,
                            default_flow_style=False )
            default_wifi = False

if default_wifi:
    print("Default wifi keys used")

# Temporary directory for the zipfile to reside in
ziptmpd = tempfile.mkdtemp( suffix="-pyenv" )
tmpzip = os.path.join( ziptmpd, "robot.zip" )

subprocess.check_call( ["zip", "-9qr", tmpzip, "./"], cwd = tmpd )

if os.path.exists( OUTPUT_ZIP ):
    os.unlink( OUTPUT_ZIP )

shutil.move( tmpzip, OUTPUT_ZIP )

shutil.rmtree( ziptmpd )
shutil.rmtree( tmpd )
