#!/usr/bin/env ruby

require 'json'
require 'faraday'
require 'faraday_middleware'

class NomadJob

  def initialize(project, image_name, nomad_branch, version, base_url, basecdn_url, nomad_addr, ecr_url, proto='https')
    validate_version version

    @project = project
    @image_name = image_name
    @nomad_branch = nomad_branch
    @base_url = base_url
    @basecdn_url = basecdn_url
    @version = version
    @ecr_url = ecr_url
    @proto = proto

    begin
      @conn = Faraday.new(url: nomad_addr) do |conn|
        conn.response :json, :content_type => 'application/json'
        conn.adapter Faraday.default_adapter
      end
    rescue
      teamcity_status 1, "Initial: Could not connect to nomad on #{nomad_addr}"
    end
  end

  def wait_for_job
    teamcity_message 'Submitting job to Nomad'
    submit_job
    teamcity_message "Nomad Job successful, waiting on evaluation #{@evaluation_id}"
    poll_evaluation
    teamcity_message "Nomad Evaluation successful, waiting on allocations #{@allocation_ids.join ','}"
    poll_allocations
    teamcity_message 'Nomad Allocations successful'
    teamcity_status 0, "Build is available at http#{@proto.eql?('https') ? 's' : ''}://#{@nomad_branch}.#{@base_url}/"
  end

  def submit_job
    resp = @conn.post do |req|
      req.url '/v1/jobs'
      req.headers['Content-Type'] = 'application/json'
      req.body = jobspec.to_json
    end

    if resp.status >= 300
      teamcity_status 1, "Got #{resp.status} from /v1/jobs: #{resp.body}"
    else
      @evaluation_id = resp.body['EvalID']
    end
  rescue
    teamcity_status 1, 'submit_job: Could not connect to nomad'
  end

  def poll_evaluation
    resp = @conn.get do |req|
      req.url "/v1/evaluation/#{@evaluation_id}/allocations"
    end

    if resp.body.length < 1
      15.times do # 30 second timeout
        resp = @conn.get do |req|
          req.url "/v1/evaluation/#{@evaluation_id}/allocations"
        end
        break if resp.body.length > 0
        sleep 2
      end

    end

    if resp.body.length < 1
      teamcity_status 1, 'Timed out receiving allocations from nomad evaluation'
    end

    @allocation_ids = resp.body.map {|r| r['ID']}
  rescue
    teamcity_status 1, 'poll_evaluation: Could not connect to nomad'
  end

  def poll_allocations

    @allocation_ids.each do |allocation|

      60.times do # 5 minute timeout

        resp = @conn.get do |req|
          req.url "/v1/allocation/#{allocation}"
        end

        unless resp.body['TaskStates'].nil?
          results = resp.body['TaskStates'].values.map {|a| a['State']}

          if results.none? {|r| r == 'pending'}
            if results.any? {|r| r == 'dead'}
              errors = resp.body['TaskStates'].values.map {|a| a['Events'].reverse.map {|e| " #{e['DriverError']}#{e['RestartReason']}"}}
              teamcity_status 1, "Nomad allocation Failed:#{errors.join(',').gsub(/'/, '"')}"
            else
              break
            end
          end
        end

        sleep 5
      end
    end

  rescue
    teamcity_status 1, 'poll_allocations: Could not connect to nomad'
  end

  def jobspec
    {
        Job: {
            ID: "#{@project}-#{@nomad_branch}",
            Name: "#{@project}-#{@nomad_branch}",
            Type: "service",
            Priority: 50,
            Datacenters: [
                "pdx-a",
                "pdx-b",
                "pdx-c",
            ],
            Constraints: [
                {
                    LTarget: "${attr.kernel.name}",
                    RTarget: "linux",
                    Operand: "="
                }
            ],
            TaskGroups: [
                {
                    Name: "#{@project}-#{@nomad_branch}",
                    Count: 2,
                    Tasks: [
                        {
                            Name: "#{@project}-#{@nomad_branch}",
                            Driver: "docker",
                            Config: {
                                image: "https://#{@ecr_url}/#{@image_name}:#{@version}",
                                port_map: [
                                    {
                                        proxy: 8080
                                    }
                                ]
                            },
                            Services: [
                                {
                                    Name: "#{@project}-#{@nomad_branch}-service",
                                    Tags: [
                                        "urlprefix-#{@nomad_branch}.#{@base_url}/",
                                        "urlprefix-#{@nomad_branch}.#{@basecdn_url}/"
                                    ],
                                    PortLabel: "proxy",
                                    Checks: [
                                        {
                                            Name: "alive",
                                            Type: "tcp",
                                            Interval: 10000000000,
                                            Timeout: 2000000000,
                                        }
                                    ]
                                }
                            ],
                            Resources: {
                                CPU: 100,
                                MemoryMB: 100,
                                Networks: [
                                    {
                                        DynamicPorts: [
                                            {
                                                Label: "proxy",
                                                Value: 0
                                            }
                                        ],
                                        MBits: 1
                                    }
                                ]
                            },
                            KillTimeout: 5000000000,
                            LogConfig: {
                                MaxFiles: 10,
                                MaxFileSizeMB: 10
                            },
                        }
                    ],
                    RestartPolicy: {
                        Interval: 300000000000,
                        Attempts: 10,
                        Delay: 25000000000,
                        Mode: "delay"
                    },
                    EphemeralDisk: {
                        SizeMB: 300
                    },
                }
            ],
            Update: {
                Stagger: 10000000000,
                MaxParallel: 1
            }
        }
    }
  end

  def validate_version version
    unless version.match(/^[a-zA-Z0-9_][-a-zA-Z0-9._]*$/) || version.length <= 128
      teamcity_status 1, "Image Version (#{version}) argument does not match docker tag syntax requirements"
    end
  end

  def teamcity_status exit_code, message
    puts "##teamcity[buildStatus status='#{exit_code == 0 ? 'Success' : 'Failed'}' text='#{message}']"
    exit exit_code
  end

  def teamcity_message message, name=nil
    puts "##teamcity[#{name || 'message'} text='#{message}']"
  end

end

class String
  def is_i?
    /\A\d+\z/ === self
  end
end

ENV['PROTO'] = 'https' unless ENV['PROTO']
job_runner = NomadJob.new(ENV['TEAMCITY_PROJECT_NAME'], ENV['IMAGE_NAME'], ENV['NOMAD_BRANCH'], ENV['IMAGE_VERSION'], ENV['BASE_URL'], ENV['BASECDN_URL'], ENV['NOMAD_ADDR'], ENV['AWS_ECR_URL'], ENV['PROTO'])
job_runner.wait_for_job

