#!/usr/local/sbin/charm-env python3

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys

from path import Path
from time import time
import subprocess

from charmhelpers.core import hookenv
from charms.benchmark import Benchmark
from charms.reactive import is_state


def fail(msg):
    hookenv.action_set({'outcome': 'failure'})
    hookenv.action_fail(msg)
    sys.exit()


def main():
    bench = Benchmark()

    if not is_state('spark.started'):
        msg = 'Spark is not started yet'
        fail(msg)

    num_partitions = hookenv.action_get('partitions') or ''

    # create dir to store results
    run = int(time())
    result_dir = Path('/opt/sparkpi-results')
    result_log = result_dir / '{}.log'.format(run)
    if not result_dir.exists():
        result_dir.mkdir()
    result_dir.chown('ubuntu', 'ubuntu')

    bench.start()
    start = int(time())

    hookenv.log("values: {} {}".format(num_partitions, result_log))

    print('calculating pi')

    # get the examples jar
    spark_home = "/usr/lib/spark"
    example_jar_name = "spark-examples.jar"
    example_jar_path = None
    for root, dirs, files in os.walk(spark_home):
        if example_jar_name in files:
            example_jar_path = os.path.join(root, example_jar_name)

    if not example_jar_path:
        msg = 'Could not find {}'.format(example_jar_name)
        fail(msg)

    with open(result_log, 'w') as log_file:
        arg_list = [
            'spark-submit',
            '--class',
            'org.apache.spark.examples.SparkPi',
            example_jar_path
        ]
        if num_partitions:
            # This is always blank. TODO: figure out what it was
            # supposed to do.
            arg_list.append(num_partitions)

        try:
            subprocess.check_call(arg_list, stdout=log_file,
                                  stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as e:
            print('smoke test command failed: ')
            print('{}'.format(' '.join(arg_list)))
            fail('spark-submit failed: {}'.format(e))

    stop = int(time())
    bench.finish()

    duration = stop - start
    bench.set_composite_score(duration, units='secs', direction='asc')

    # Tell the user how they can get the full result log
    scp_log_msg = "juju scp {}:{} .".format(hookenv.local_unit(), result_log)
    hookenv.action_set({'meta.fetchcmd': scp_log_msg})

    with open(result_log) as log:
        success = False
        for line in log.readlines():
            if 'Pi is roughly 3.1' in line:
                success = True
                pi_calc = line
                break

    if not success:
        msg = 'Spark-submit failed to calculate pi'
        fail(msg)

    hookenv.action_set({'meta.raw': pi_calc})
    hookenv.action_set({'outcome': 'success'})


if __name__ == '__main__':
    main()
