#!/usr/bin/env python

import argparse
import os
from docker import Client
from functools import partial
from pprint import pprint

from mexec.mesos import Mesos
from mexec.docker import Docker


def set_mesos_host(hosts):
    hosts_str = hosts or os.environ.get('MEXEC_MESOS_MASTER', '')
    if not hosts_str:
        raise Exception('Argument --mesos-master is required or ' +
                        'MEXEC_MESOS_MASTER env variable is have to be set!')
    return hosts_str.split(',') if hosts_str else ''


def set_docker_port(port):
    port = port or os.environ.get('MEXEC_DOCKER_PORT', '')
    if not port:
        raise Exception('Argument --docker-port is required or ' +
                        'MEXEC_DOCKER_PORT env variable is have to be set!')
    return port


def print_results(results):
    for result in results:
        print('=' * 200)
        if isinstance(result, str) or isinstance(result, unicode):
            print(result)
        else:
            pprint(result)


def implicit_cast(value):
    if value == 'True':
        return True
    elif value == 'False':
        return False
    elif value.isdigit():
        return int(value)
    else:
        try:
            return float(value)
        except ValueError:
            return value


def parse_arg(arg):
    key, value = arg.split('=', 1)
    return (key, implicit_cast(value))


if __name__ == '__main__':
    desc = 'Execute commands inside docker containers run by marathon.'
    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument('task_id', help='Mesos task id')
    parser.add_argument('action',
                        choices=['exec_container', 'inspect', 'top', 'logs'],
                        help='Action on the container')
    parser.add_argument('--arg', dest='args', action='append', default=[],
                        type=parse_arg, help='Arguments of the docker action')
    parser.add_argument('--mesos-master', type=set_mesos_host, default='',
                        help='Mesos master host separated by commas eg: ' +
                             'master-1:5050,master-2:5050,master-3:5050')
    parser.add_argument('--docker-port', type=set_docker_port, default='',
                        help='Docker tcp port')

    args = parser.parse_args()

    mesos = Mesos(args.mesos_master)
    tasks = mesos.get_mesos_tasks(args.task_id)
    slaves = mesos.get_slave_hostnames(set([task.slave_id for task in tasks]))

    for task in tasks:
        cli = Client(base_url='tcp://{}:{}'.format(slaves[task.slave_id],
                                                   args.docker_port))
        docker = Docker(cli)
        env_var = 'MESOS_TASK_ID={}'.format(task.task_id)

        container_callback = partial(getattr(docker, args.action),
                                     **dict(args.args))
        results = docker.exec_by_env_variable(env_var, container_callback)

        print_results(results)

    if not len(tasks):
        print('No tasks found!')
