package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	docker "github.com/fsouza/go-dockerclient"
)

// OutputCollector is an io.Writer that accumulates output from a specified stream in an attached
// Docker container and appends it to the appropriate field within a SubmittedJob.
type OutputCollector struct {
	context  *Context
	job      *SubmittedJob
	isStdout bool
}

// DescribeStream returns "stdout" or "stderr" to indicate which stream this collector is consuming.
func (c OutputCollector) DescribeStream() string {
	if c.isStdout {
		return "stdout"
	}
	return "stderr"
}

// Write appends bytes to the selected stream and updates the SubmittedJob.
func (c OutputCollector) Write(p []byte) (int, error) {
	log.WithFields(log.Fields{
		"length": len(p),
		"bytes":  string(p),
		"stream": c.DescribeStream(),
	}).Debug("Received output from a job")

	if c.isStdout {
		c.job.Stdout += string(p)
	} else {
		c.job.Stderr += string(p)
	}

	if err := c.context.UpdateJob(c.job); err != nil {
		return 0, err
	}

	return len(p), nil
}

// Runner is the main entry point for the job runner goroutine.
func Runner(c *Context) {
	for {
		Claim(c)

		time.Sleep(time.Duration(c.Poll) * time.Millisecond)
	}
}

// Claim acquires the oldest single pending job and launches a goroutine to execute its command in
// a new container.
func Claim(c *Context) {
	job, err := c.ClaimJob()
	if err != nil {
		log.WithFields(log.Fields{"error": err}).Error("Unable to claim a job.")
		return
	}
	if job == nil {
		// Nothing to claim.
		return
	}
	if err := job.Validate(); err != nil {
		fields := log.Fields{
			"jid":     job.JID,
			"account": job.Account,
			"error":   err,
		}

		log.WithFields(fields).Error("Invalid job in queue.")

		job.Status = StatusError
		if err := c.UpdateJob(job); err != nil {
			fields["error"] = err
			log.WithFields(fields).Error("Unable to update job status.")
		}

		return
	}

	go Execute(c, job)
}

// Execute launches a container to process the submitted job. It passes any provided stdin data
// to the container and consumes stdout and stderr, updating Mongo as it runs. Once completed, it
// acquires the job's result from its configured source and marks the job as finished.
func Execute(c *Context, job *SubmittedJob) {
	defaultFields := log.Fields{
		"jid":     job.JID,
		"account": job.Account,
	}

	// Logging utility messages.
	debug := func(message string) {
		log.WithFields(defaultFields).Debug(message)
	}
	reportErr := func(message string, err error) {
		fs := log.Fields{}
		for k, v := range defaultFields {
			fs[k] = v
		}
		fs["err"] = err
		log.WithFields(fs).Error(message)
	}
	checkErr := func(message string, err error) bool {
		if err == nil {
			debug(fmt.Sprintf("%s: ok", message))
			return false
		}

		reportErr(fmt.Sprintf("%s: ERROR", message), err)
		return true
	}

	// Update the job model in Mongo, reporting any errors along the way.
	// This also updates our job model with any changes from Mongo, such as the kill request flag.
	updateJob := func(message string) bool {
		if err := c.UpdateJob(job); err != nil {
			reportErr(fmt.Sprintf("Unable to update the job's %s.", message), err)
			return false
		}
		return true
	}

	log.WithFields(defaultFields).Info("Launching a job.")

	job.StartedAt = StoreTime(time.Now())
	job.QueueDelay = job.StartedAt.AsTime().Sub(job.CreatedAt.AsTime()).Nanoseconds()

	image := c.DefaultImage
	if len(job.Layers) != 0 {
		image = job.Layers[0].Name
		if len(job.Layers) > 1 {
			log.WithFields(log.Fields{
				"layers": job.Layers,
			}).Warn("More than one layer used")
		}

	}

	container, err := c.CreateContainer(docker.CreateContainerOptions{
		Name: job.ContainerName(),
		Config: &docker.Config{
			Image:     image,
			Cmd:       []string{"/bin/bash", "-c", job.Command},
			OpenStdin: true,
			StdinOnce: true,
		},
	})
	if checkErr("Created the job's container", err) {
		job.Status = StatusError
		updateJob("status")
		return
	}

	// Record the job's container ID.
	job.ContainerID = container.ID
	if !updateJob("start timestamp and container id") {
		return
	}

	// Include container information in this job's logging messages.
	defaultFields["container id"] = container.ID
	defaultFields["container name"] = container.Name

	// Was a kill requested between the time the job was claimed, and the time the container was
	// created? If so: transition the job to StatusKilled and jump ahead to removing the container
	// we just created. If not: continue with job execution normally.

	// If a kill was requested before the job was claimed, it would have been removed from the queue.
	// If a kill is requested after the container was created, it will have the containerID that we
	// just sent and be able to kill the running container.

	if job.KillRequested {
		job.Status = StatusKilled
	} else {
		// Prepare the input and output streams.
		stdin := bytes.NewReader(job.Stdin)
		stdout := OutputCollector{
			context:  c,
			job:      job,
			isStdout: true,
		}
		stderr := OutputCollector{
			context:  c,
			job:      job,
			isStdout: false,
		}

		go func() {
			err = c.AttachToContainer(docker.AttachToContainerOptions{
				Container:    container.ID,
				Stream:       true,
				InputStream:  stdin,
				OutputStream: stdout,
				ErrorStream:  stderr,
				Stdin:        true,
				Stdout:       true,
				Stderr:       true,
			})
			checkErr("Attached to the container", err)
		}()

		// Start the created container.
		err = c.StartContainer(container.ID, &docker.HostConfig{})
		if checkErr("Started the container", err) {
			job.Status = StatusError
			updateJob("status")
			return
		}

		// Measure the container-launch overhead here.
		overhead := time.Now()
		job.OverheadDelay = overhead.Sub(job.StartedAt.AsTime()).Nanoseconds()
		updateJob("overhead delay")

		status, err := c.WaitContainer(container.ID)
		if checkErr("Waited for the container to complete", err) {
			job.Status = StatusError
			updateJob("status")
			return
		}

		job.FinishedAt = StoreTime(time.Now())
		job.Runtime = job.FinishedAt.AsTime().Sub(overhead).Nanoseconds()
		if status == 0 {
			// Successful termination.
			job.Status = StatusDone

			// Extract the result from the job.
			if job.ResultSource == "stdout" {
				job.Result = []byte(job.Stdout)
				debug("Acquired job result from stdout: ok")
			} else if strings.HasPrefix(job.ResultSource, "file:") {
				resultPath := job.ResultSource[len("file:"):len(job.ResultSource)]

				var resultBuffer bytes.Buffer
				err = c.CopyFromContainer(docker.CopyFromContainerOptions{
					Container:    container.ID,
					Resource:     resultPath,
					OutputStream: &resultBuffer,
				})
				if checkErr(fmt.Sprintf("Acquired the job's result from the file [%s]", resultPath), err) {
					job.Status = StatusError
				} else {
					// CopyFromContainer returns the file contents as a tarball.
					var content bytes.Buffer
					r := bytes.NewReader(resultBuffer.Bytes())
					tr := tar.NewReader(r)

					for {
						_, err := tr.Next()
						if err == io.EOF {
							break
						}
						if err != nil {
							reportErr("Read tar-encoded content: ERROR", err)
							job.Status = StatusError
							break
						}

						if _, err = io.Copy(&content, tr); err != nil {
							reportErr("Copy decoded content: ERROR", err)
							job.Status = StatusError
							break
						}
					}

					job.Result = content.Bytes()
				}
			}
		} else {
			// Something went wrong.

			// See if a kill was explicitly requested. If so, transition to StatusKilled. Otherwise,
			// transition to StatusError.
			killed, err := c.JobKillRequested(job.JID)
			if err != nil {
				reportErr("Check the job kill status: ERROR", err)
				return
			}

			if killed {
				job.Status = StatusKilled
			} else {
				job.Status = StatusError
			}
		}

		// Job execution has completed successfully.
	}

	err = c.RemoveContainer(docker.RemoveContainerOptions{ID: container.ID})
	checkErr("Removed the container", err)

	err = c.UpdateAccountUsage(job.Account, job.Runtime)
	if err != nil {
		reportErr("Update account usage: ERROR", err)
	}
	updateJob("status and final result")

	log.WithFields(log.Fields{
		"jid":      job.JID,
		"account":  job.Account,
		"status":   job.Status,
		"runtime":  job.Runtime,
		"overhead": job.OverheadDelay,
		"queue":    job.QueueDelay,
	}).Info("Job complete.")
}
