package main

// Copyright 2016-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance with the License. A copy of the License is located at
// http://aws.amazon.com/apache2.0/
// or in the "license" file accompanying this file. This file 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 (
    "fmt"
    log "github.com/Sirupsen/logrus"
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/ec2metadata"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/route53"
    "github.com/fsouza/go-dockerclient"
    "github.com/aws/aws-sdk-go/service/cloudwatchevents"
    "net/http"
    "io/ioutil"
    "strconv"
    "strings"
    "time"
    "flag"
)

const workerTimeout = 180 * time.Second
const defaultTTL = 0
const defaultWeight = 1

var DNSName = "alissa.agilent.com"

type handler interface {
    Handle(*docker.APIEvents) error
}

type dockerRouter struct {
    handlers      map[string][]handler
    dockerClient  *docker.Client
    listener      chan *docker.APIEvents
    workers       chan *worker
    workerTimeout time.Duration
}

func dockerEventsRouter(bufferSize int, workerPoolSize int, dockerClient *docker.Client,
    handlers map[string][]handler) (*dockerRouter, error) {
    workers := make(chan *worker, workerPoolSize)
    for i := 0; i < workerPoolSize; i++ {
        workers <- &worker{}
    }

    dockerRouter := &dockerRouter{
        handlers:      handlers,
        dockerClient:  dockerClient,
        listener:      make(chan *docker.APIEvents, bufferSize),
        workers:       workers,
        workerTimeout: workerTimeout,
    }

    return dockerRouter, nil
}

func (e *dockerRouter) start() error {
    go e.manageEvents()
    return e.dockerClient.AddEventListener(e.listener)
}

func (e *dockerRouter) stop() error {
    if e.listener == nil {
        return nil
    }
    return e.dockerClient.RemoveEventListener(e.listener)
}

func (e *dockerRouter) manageEvents() {
    for {
        event := <-e.listener
        timer := time.NewTimer(e.workerTimeout)
        gotWorker := false
        // Wait until we get a free worker or a timeout
        // there is a limit in the number of concurrent events managed by workers to avoid resource exhaustion
        // so we wait until we have a free worker or a timeout occurs
        for !gotWorker {
            select {
            case w := <-e.workers:
                if !timer.Stop() {
                    <-timer.C
                }
                go w.doWork(event, e)
                gotWorker = true
            case <-timer.C:
                log.Infof("Timed out waiting.")
            }
        }
    }
}

type worker struct{}

func (w *worker) doWork(event *docker.APIEvents, e *dockerRouter) {
    defer func() { e.workers <- w }()
    if handlers, ok := e.handlers[event.Status]; ok {
        log.Infof("Processing event: %#v", event)
        for _, handler := range handlers {
            if err := handler.Handle(event); err != nil {
                log.Errorf("Error processing event %#v. Error: %v", event, err)
            }
        }
    }
}

type dockerHandler struct {
    handlerFunc func(event *docker.APIEvents) error
}

func (th *dockerHandler) Handle(event *docker.APIEvents) error {
    return th.handlerFunc(event)
}

type config struct {
    HostedZoneId string
    Hostname     string
    Region       string
}

var configuration config

func logErrorAndFail(err error) {
    if err != nil {
        log.Fatal(err)
    }
}

func logErrorNoFatal(err error) {
    if err != nil {
        log.Error(err)
    }
}

type topTasks struct {
    Tasks []taskInfo
}

type taskInfo struct {
    Arn           string
    DesiredStatus string
    KnownStatus   string
    Family        string
    Version       string
    Containers    []ContainerInfo
}

type ContainerInfo struct {
    DockerId   string
    DockerName string
    Name       string
}

type ServiceInfo struct {
    Name string
    Port string
}

func getDNSHostedZoneId() (string, error) {
    r53 := route53.New(session.New())
    params := &route53.ListHostedZonesByNameInput{
        DNSName: aws.String(DNSName),
    }

    zones, err := r53.ListHostedZonesByName(params)

    if err == nil {
        if len(zones.HostedZones) > 0 {
            return aws.StringValue(zones.HostedZones[0].Id), nil
        }
    }

    return "", err
}

func createDNSRecord(serviceName string, dockerId string, port string) error {
    r53 := route53.New(session.New())
    srvRecordName := serviceName + "." + DNSName
    // This API call creates a new DNS record for this service
    params := &route53.ChangeResourceRecordSetsInput{
        ChangeBatch: &route53.ChangeBatch{
            Changes: []*route53.Change{
                {
                    Action: aws.String(route53.ChangeActionCreate),
                    ResourceRecordSet: &route53.ResourceRecordSet{
                        Name: aws.String(srvRecordName),
                        // It creates a SRV record with the name of the service
                        Type: aws.String(route53.RRTypeSrv),
                        ResourceRecords: []*route53.ResourceRecord{
                            {
                                // priority: the priority of the target host, lower value means more preferred
                                // weight: A relative weight for records with the same priority, higher value means more preferred
                                // port: the TCP or UDP port on which the service is to be found
                                // target: the canonical hostname of the machine providing the service
                                Value: aws.String("1 1 " + port + " " + configuration.Hostname),
                            },
                        },
                        SetIdentifier: aws.String(dockerId),
                        // TTL=0 to avoid DNS caches
                        TTL:    aws.Int64(defaultTTL),
                        Weight: aws.Int64(defaultWeight),
                    },
                },
            },
            Comment: aws.String("Service Discovery Created Record"),
        },
        HostedZoneId: aws.String(configuration.HostedZoneId),
    }
    _, err := r53.ChangeResourceRecordSets(params)
    logErrorNoFatal(err)
    fmt.Println("Record " + srvRecordName + " created (1 1 " + port + " " + configuration.Hostname + ")")
    return err
}

func deleteDNSRecord(serviceName string, dockerId string) error {
    var err error
    r53 := route53.New(session.New())
    srvRecordName := serviceName + "." + DNSName
    // This API Call looks for the Route53 DNS record for this service and docker ID to get the values to delete
    paramsList := &route53.ListResourceRecordSetsInput{
        HostedZoneId:          aws.String(configuration.HostedZoneId), // Required
        MaxItems:              aws.String("10"),
        StartRecordIdentifier: aws.String(dockerId),
        StartRecordName:       aws.String(srvRecordName),
        StartRecordType:       aws.String(route53.RRTypeSrv),
    }
    resp, err := r53.ListResourceRecordSets(paramsList)
    logErrorNoFatal(err)
    if err != nil {
        return err
    }
    srvValue := ""
    for _, rrset := range resp.ResourceRecordSets {
        if *rrset.SetIdentifier == dockerId && (*rrset.Name == srvRecordName || *rrset.Name == srvRecordName+".") {
            for _, rrecords := range rrset.ResourceRecords {
                srvValue = aws.StringValue(rrecords.Value)
                break
            }
        }
    }
    if srvValue == "" {
        log.Error("Route53 Record doesn't exist")
        return nil
    }

    // This API call deletes the DNS record for the service for this docker ID
    params := &route53.ChangeResourceRecordSetsInput{
        ChangeBatch: &route53.ChangeBatch{
            Changes: []*route53.Change{
                {
                    Action: aws.String(route53.ChangeActionDelete),
                    ResourceRecordSet: &route53.ResourceRecordSet{
                        Name: aws.String(srvRecordName),
                        Type: aws.String(route53.RRTypeSrv),
                        ResourceRecords: []*route53.ResourceRecord{
                            {
                                Value: aws.String(srvValue),
                            },
                        },
                        SetIdentifier: aws.String(dockerId),
                        TTL:           aws.Int64(defaultTTL),
                        Weight:        aws.Int64(defaultWeight),
                    },
                },
            },
        },
        HostedZoneId: aws.String(configuration.HostedZoneId),
    }
    _, err = r53.ChangeResourceRecordSets(params)
    logErrorNoFatal(err)
    fmt.Println("Record " + srvRecordName + " deleted ( " + srvValue + ")")
    return err
}

var dockerClient *docker.Client

func getNetworkPortAndServiceName(container *docker.Container, includePort bool) []ServiceInfo {
    // One of the environment varialbles should be SERVICE_<port>_NAME = <name of the service>
    // We look for this environment variable doing a split in the "=" and another one in the "_"
    // So envEval = [SERVICE_<port>_NAME, <name>]
    // nameEval = [SERVICE, <port>, NAME]
    var svc []ServiceInfo = make([]ServiceInfo, 0)
    for _, env := range container.Config.Env {
        envEval := strings.Split(env, "=")
        nameEval := strings.Split(envEval[0], "_")
        if len(envEval) == 2 && len(nameEval) == 3 && nameEval[0] == "SERVICE" && nameEval[2] == "NAME" {
            if _, err := strconv.Atoi(nameEval[1]); err == nil {
                if includePort {
                    for srcPort, mapping := range container.NetworkSettings.Ports {
                        portEval := strings.Split(string(srcPort), "/")
                        if len(portEval) > 0 && portEval[0] == nameEval[1] {
                            if len(mapping) > 0 {
                                svc = append(svc, ServiceInfo{envEval[1], mapping[0].HostPort})
                            }
                        }
                    }
                } else {
                    svc = append(svc, ServiceInfo{envEval[1], ""})
                }
            }
        }
    }
    return svc
}

func sendToCWEvents (detail string, detailType string, resource string, source string) error {
    config := aws.NewConfig().WithRegion(configuration.Region)
    sess := session.New(config)
    svc := cloudwatchevents.New(sess)
    params := &cloudwatchevents.PutEventsInput{
        Entries: []*cloudwatchevents.PutEventsRequestEntry{
            {
                Detail: aws.String(detail),
                DetailType: aws.String(detailType),
                Resources: []*string{
                    aws.String(resource),
                },
                Source: aws.String(source),
                Time: aws.Time(time.Now()),
            },
        },
    }
    _, err := svc.PutEvents(params)
    logErrorNoFatal(err)
    return err
}

func getTaskArn(dockerID string) string {
    resp, err := http.Get("http://127.0.0.1:51678/v1/tasks")
    if err != nil {
        logErrorAndFail(err)
    }
    defer resp.Body.Close()
    body, err := ioutil.ReadAll(resp.Body)
    bodyStr := string(body)
    idIndex := strings.Index(bodyStr, string(dockerID))
    arnStartIndex := strings.LastIndex(bodyStr[:idIndex], "arn:aws:ecs:")
    arnString := bodyStr[arnStartIndex:]
    arnEndIndex := strings.Index(arnString, "\"")
    return arnString[:arnEndIndex]
}

func main() {
    var err error
    var sum int
    var zoneId string

    var sendEvents = flag.Bool("cw-send-events", false, "Send CloudWatch events when a container is created or terminated")
    var zoneIdStr = flag.String("hosted-zone-id", "none", "Route53 Hosted Zone Id")

    flag.Parse()

    var DNSNameArg = flag.Arg(0)
    if DNSNameArg != "" {
        DNSName = DNSNameArg
    }

    if zoneIdStr != nil {
        zoneId = *zoneIdStr
    } else {
        for {
            // We try to get the Hosted Zone Id using exponential backoff
            zoneId, err = getDNSHostedZoneId()
            if err == nil {
                break
            }
            if sum > 8 {
                logErrorAndFail(err)
            }
            time.Sleep(time.Duration(sum) * time.Second)
            sum += 2
        }
    }
    configuration.HostedZoneId = zoneId
    metadataClient := ec2metadata.New(session.New())
    hostname, err := metadataClient.GetMetadata("/hostname")
    configuration.Hostname = hostname
    logErrorAndFail(err)
    region, err := metadataClient.Region()
    configuration.Region = region
    logErrorAndFail(err)

    endpoint := "unix:///var/run/docker.sock"
    startFn := func(event *docker.APIEvents) error {
        var err error
        container, err := dockerClient.InspectContainer(event.ID)
        logErrorAndFail(err)
        allService := getNetworkPortAndServiceName(container, true)
        for _, svc := range allService {
            if svc.Name != "" && svc.Port != "" {
                sum = 1
                for {
                    if err = createDNSRecord(svc.Name, event.ID, svc.Port); err == nil {
                        break
                    }
                    if sum > 8 {
                        log.Error("Error creating DNS record")
                        break
                    }
                    time.Sleep(time.Duration(sum) * time.Second)
                    sum += 2
                }
            }
        }
        if *sendEvents {
            taskArn := getTaskArn(event.ID)
            sendToCWEvents(`{ "dockerId": "` + event.ID + `","TaskArn":"` + taskArn + `" }`, "Task Started", configuration.Hostname, "awslabs.ecs.container" )
        }
        fmt.Println("Docker " + event.ID + " started")
        return nil
    }

    stopFn := func(event *docker.APIEvents) error {
        var err error
        container, err := dockerClient.InspectContainer(event.ID)
        logErrorAndFail(err)
        allService := getNetworkPortAndServiceName(container, false)
        for _, svc := range allService {
            if svc.Name != "" {
                sum = 1
                for {
                    if err = deleteDNSRecord(svc.Name, event.ID); err == nil {
                        break
                    }
                    if sum > 8 {
                        log.Error("Error deleting DNS record")
                        break
                    }
                    time.Sleep(time.Duration(sum) * time.Second)
                    sum += 2
                }
            }
        }
        if *sendEvents {
            taskArn := getTaskArn(event.ID)
            sendToCWEvents(`{ "dockerId": "` + event.ID + `","TaskArn":"` + taskArn + `" }`, "Task Stopped", configuration.Hostname, "awslabs.ecs.container" )
        }
        fmt.Println("Docker " + event.ID + " stopped")
        return nil
    }

    startHandler := &dockerHandler{
        handlerFunc: startFn,
    }
    stopHandler := &dockerHandler{
        handlerFunc: stopFn,
    }
    handlers := map[string][]handler{"start": []handler{startHandler}, "die": []handler{stopHandler}}

    dockerClient, _ = docker.NewClient(endpoint)
    router, err := dockerEventsRouter(5, 5, dockerClient, handlers)
    logErrorAndFail(err)
    defer router.stop()
    router.start()
    fmt.Println("Waiting events")
    select {}
}