#!/usr/bin/env python
import time

import a_sync
import simplejson as json

from paasta_tools import mesos_tools
from paasta_tools.mesos.exceptions import SlaveDoesNotExist
from paasta_tools.utils import load_system_paasta_config
from paasta_tools.utils import paasta_print


def get_container_info_from_task(task):
    for status in task["statuses"]:
        if status["state"] != "TASK_RUNNING":
            continue
        container_id = (
            status.get("container_status", {}).get("container_id", {}).get("value")
        )
        time_start = status.get("timestamp")
        return container_id, time_start
    return None, None


def get_paasta_service_instance_from_task(task):
    try:
        docker_params = task["container"].get("docker", {}).get("parameters", [])
    except KeyError:
        return None, None
    service, instance = None, None
    for param in docker_params:
        if param["key"] == "label":
            label = param["value"]
            if label.startswith("paasta_service="):
                service = label.split("=")[1]
            if label.startswith("paasta_instance="):
                instance = label.split("=")[1]
    return service, instance


async def get_pool_from_task(task):
    try:
        attributes = (await task.slave())["attributes"]
        return attributes.get("pool", "default")
    except SlaveDoesNotExist:
        return None


@a_sync.to_blocking
async def get_task_allocation_info():
    tasks = await mesos_tools.get_cached_list_of_running_tasks_from_frameworks()
    info_list = []
    for task in tasks:
        info = {"resources": task["resources"]}
        info["mesos_container_id"], info["start_time"] = get_container_info_from_task(
            task
        )
        info["paasta_service"], info[
            "paasta_instance"
        ] = get_paasta_service_instance_from_task(task)
        info["paasta_pool"] = await get_pool_from_task(task)
        info_list.append(info)
    return info_list


def main():
    cluster = load_system_paasta_config().get_cluster()
    info_list = get_task_allocation_info()
    timestamp = time.time()
    for info in info_list:
        info["cluster"] = cluster
        info["timestamp"] = timestamp
        paasta_print(json.dumps(info))


if __name__ == "__main__":
    main()
