#!/usr/bin/env python3
import npyscreen,subprocess,json,os,sys,boto3,time,copy,argparse,re
from threading import Thread
from itertools import chain,islice
import webbrowser

def ip(item):
    npyscreen.notify("Resolving TaskArn", title='Please Wait', form_color='DEFAULT', wide=True)
    if "taskarn" in item:
        tasks = client.describe_tasks(cluster = item["cluster"], tasks = [item["taskarn"]])['tasks']
        item["containerInstance"] = tasks[0]['containerInstanceArn']
        item["hostPort"] = [i for i in tasks[0]['containers'] if i['containerArn']==item['containerarn']][0]['networkBindings'][0]['hostPort']
        npyscreen.notify("Resolving ContainerInstance", title='Please Wait', form_color='DEFAULT', wide=True)
    if "containerInstance" in item:
        item["instance"] = client.describe_container_instances(cluster = item["cluster"], containerInstances = [item["containerInstance"]])['containerInstances'][0]['ec2InstanceId']
    npyscreen.notify("Resolving PrivateIpAddress", title='Please Wait', form_color='DEFAULT', wide=True)
    if "instance" in item:
        return str(ec2.describe_instances(InstanceIds=[item["instance"]])['Reservations'][0]['Instances'][0]['PrivateIpAddress'])
    else:
        return str(item)

def container(item):
    if "taskarn" in item:
        return item["taskarn"]
    else:
        return None

def port(item):
    if "hostPort" in item:
        return item["hostPort"]
    else:
        return None

class ContainerTree(npyscreen.MLTree):
    def handle_mouse_event(self,m):
        self.never_cache=True
        super(ContainerTree,self).handle_mouse_event(m)
        self.h_select(m)
        
    def h_select(self,ch):
        if self.cursor_line==0:
            return
        item = self.values[self.cursor_line]
        if not item.address:
            return
        if args.browser:
           ipa = ip(item.address)
           porta = port(item.address)
           webbrowser.open_new("http://"+ipa+":"+str(porta))
           return
        global execl
        if os.getenv('SSH_OPTS'):
            ssh_opts = [i for i in os.getenv('SSH_OPTS', "").split(" ") if len(i)>0]
        else:
            ssh_opts = ["-t","-A","-o","StrictHostKeyChecking=no","-l","ec2-user"]
        i=ip(item.address)
        cont=container(item.address)
        if cont:
            execl = ["/usr/bin/ssh"]+ssh_opts+[ i, "docker exec -it $(docker ps -q -f label=com.amazonaws.ecs.task-arn="+cont+") /bin/sh"]
        else:
            execl = ["/usr/bin/ssh"]+ssh_opts+[ i ]
        self.parentForm.parentApp.switchForm(None)

class BaseTreeForm(npyscreen.FormBaseNew):
    def create(self):
        services_first = False
        title="ECS"
        treedata = npyscreen.TreeData(content=title, selectable=True, selected=True, ignore_root=False)
        self.treedata = treedata
        self.add_event_hander("INSERT", self.ev_insert_handler)
        self.add_event_hander("DELETE", self.ev_delete_handler)
        self.add_event_hander("UPDATE", self.ev_update_handler)
        wgtree = self.add(ContainerTree)
        self.wgtree=wgtree
        wgtree.parentForm = self
        wgtree.values = treedata

    def insert(self, tree, hier, address):
        if len(hier) > 0:
            value = hier.pop(0)
            children = []
            for i in tree._children:
                if i.content == value:
                    self.insert(i, hier, address)
                    break
            else:
                child = tree.new_child(content = value, selectable = True)
                child.operation = False
                child.address = address
                self.insert(child, hier, address)

    def update(self, tree, hier, address, replace):
        if len(hier) > 0:
            value = hier.pop(0)
            children = []
            for i in tree._children:
                if i.content == value:
                    if len(hier) == 0:
                        i.operation = False
                        i.content = replace
                        i.address = address
                    else:
                        self.update(i, hier, address, replace)
                    break

    def ev_insert_handler(self, event):
        hier, operation, address = event.payload
        self.insert(self.treedata, hier, address)
        self.wgtree.resize()

    def ev_delete_handler(self, event):
        pass

    def ev_update_handler(self, event):
        hier, operation, address, replace = event.payload
        self.update(self.treedata, hier, address, replace)
        self.wgtree.resize()


class ContainerTreeApp(npyscreen.StandardApp):
    def onStart(self):
        self.addForm("MAIN", BaseTreeForm, framed=False, max_height=15)

    def onCleanExit(self):
        pass

    def change_form(self, name):
        self.switchForm(name)

def cl(i):
    return i.split("/",1)[1]

def boit(arnname,command,**args):
    return chain.from_iterable([i[arnname] for i in client.get_paginator(command).paginate(**args)])
        
class EventSender(Thread):
    def __init__(self,seto):
        super(EventSender, self).__init__()
        self.seto = seto
        self.daemon = True

    def insert(self,*data):
        self.seto.queue_event(npyscreen.Event("INSERT",copy.deepcopy(data)))

    def update(self,*data):
        self.seto.queue_event(npyscreen.Event("UPDATE",copy.deepcopy(data)))
        
    def run(self):
        al = []
        vers = {}
        inst = {}
        for c in boit('clusterArns','list_clusters'):
            if args.cluster and re.search(args.cluster,cl(c)) is None:
                continue
            al.append(cl(c))
            self.insert(al,False,False)
            inst[c]=[]
            vers[c]={}
            if args.ec2:
                for b in boit('serviceArns','list_services',cluster=c):                    
                    al.append(cl(b))
                    self.insert(al,False,False)
                    for l in boit('taskArns','list_tasks',cluster=c,serviceName=b):
                        al.append(cl(l))
                        vers[c][l]=copy.deepcopy(al)
                        self.insert(al,False,{"cluster":c,"taskarn":l})
                        al.pop()
                    al.pop()
            else:
                for b in boit('containerInstanceArns','list_container_instances',cluster=c):
                    al.append(cl(b))
                    inst[c].append(b)
                    self.insert(al,False,{"cluster":c,"containerInstance":b})
                    for l in boit('taskArns','list_tasks',cluster=c,containerInstance=b):
                        al.append(cl(l))
                        vers[c][l]=copy.deepcopy(al)
                        self.insert(al,False,{"cluster":c,"containerInstance":b})
                        al.pop()
                    al.pop()
            al.pop()
        for c,c_v in vers.items():
            for task in c_v.keys():
                for t in client.describe_tasks(cluster = c, tasks = [task])['tasks']:
                    self.update(vers[c][t['taskArn']],False,{"cluster":c,"containerInstance":t['containerInstanceArn']},cl(t['taskDefinitionArn']))
                    for i, ct in enumerate(t['containers']):
                        taddr = vers[c][t['taskArn']]
                        if args.no_show_tasks:                            
                            if i==0:
                                self.update(taddr[:-1]+[cl(task)], False, {"cluster":c,"containerInstance":t['containerInstanceArn'], "taskarn":ct['taskArn'], "containerarn": ct['containerArn']},ct['name'])
                                continue
                            taddr = taddr[:-1]+[ct['name']]
                        else:
                            taddr = taddr+[ct['name']]
                        self.insert(taddr, False, {"cluster":c,"containerInstance":t['containerInstanceArn'], "taskarn":ct['taskArn'], "containerarn": ct['containerArn']})
        if not args.ec2:
            for c,c_i in inst.items():
                for i in c_i:
                    for di in client.describe_container_instances(cluster = c, containerInstances = [i])['containerInstances']:
                        self.update([cl(c),cl(i)],False,{"instance":di['ec2InstanceId']},di['ec2InstanceId'])




def main():
   try:
       parser = argparse.ArgumentParser()
       parser.add_argument("-C", "--onlyclusters",help="show clusters only", action="store_true")
       parser.add_argument("-c", "--cluster", type=str, help="show specific cluster only")
       parser.add_argument("-r", "--region", type=str, help="aws region")
       parser.add_argument("-p", "--profile", type=str, help="aws profile")
       parser.add_argument("-e", "--ec2", help="show ec2 instances", action="store_false")
       parser.add_argument("-b", "--browser", help="open browser", action="store_true")
       parser.add_argument("-t", "--no_show_tasks", help="don't show tasks in branches (on by default)", action="store_false")
       parser.add_argument("-s", "--no_show_services", help="don't show services in clusters", action="store_true")
       parser.add_argument("-i", "--no_show_instances", help="don't show instances in clusters", action="store_true")

       global args
       global client
       global ec2
       args = parser.parse_args()
       session = {}
       if args.profile:
           session["profile_name"] = args.profile 
       if args.region:
           session["region_name"] = args.region
       if len(session) > 0:
           boto3.setup_default_session(**session)
       client = boto3.client('ecs')
       ec2    = boto3.client('ec2')
       TA = ContainerTreeApp()
       EV = EventSender(TA)
       EV.start()
       TA.run()
       if len(execl) > 0:
           print(" ".join([str(i) for i in execl]))
           sys.stdout.flush()
           subprocess.call(execl)
   except KeyboardInterrupt:
       print("Ctrl-c")


if __name__ == '__main__':
    main()
