package main

import (
    "gopkg.in/cheggaaa/pb.v1"
	"bytes"
	"bufio"
	"flag"
	"fmt"
	"image"
	"log"
	"net"
	"os"
	"time"
     _ "image/jpeg"
)

var (

    // Define command line arguments
	host = flag.String("host", "127.0.0.1", "Server address (v4)")
	port = flag.String("port", "1337", "Server port")
	address string

	imagePath = flag.String("image", "", "Relative image path")
    xoffset = flag.Int("xoffset", 0, "xoffset")
	yoffset = flag.Int("yoffset", 0, "yoffset")

	canvasWidth = flag.Int("width", 1920, "canvas width")
	canvasHeight = flag.Int("height", 1080, "canvas height")

	testConn = flag.Bool("t", false, "test the connection before escalating completely")

	fill = flag.Bool("fill", false, "fill the complete canvas")
	color = flag.String("col", "000000", "define a color")

	cores = flag.Int("cores", 1, "Amount of cores to use")
	loops = flag.Int("loops", 500, "Amount os loops to run")
)

// parse the command line
func parseFlags() {
	flag.Parse()

	address = fmt.Sprintf("%s:%s", *host, *port)
}

// testConnection to the given server if the -t flag was set
func testConnection() {
	if *testConn == true {
		log.Printf("[ ] Testing Connection: True")

		log.Println("[ ] Testing TCP Connection...")
		testConnectionProtocol("tcp")

	} else {
		log.Printf("[ ] Testing Connection: False")
	}
}

// testConnectionProtocol tests the connection to the server using the given protocol ("udp" or "tcp")
func testConnectionProtocol(protocol string) {
    log.Printf("[D] address: %s", address)
    connection, netDialErr := net.Dial(protocol, address)
	if netDialErr != nil {
		log.Fatal(netDialErr)
	}
	connectionCloseError := connection.Close()
	if connectionCloseError != nil {
		log.Fatal(connectionCloseError)
	}
}

// buildSendString builds a string representing the rectangle defined in fill mode
func buildSendString(start int, end int, doneChannel chan bool, stripBufferChannel chan bytes.Buffer) {
	var stripBuffer bytes.Buffer
	var currentCommand string

	for x := start; x < end; x++ {
		for y := 0; y < *canvasHeight; y++ {
			// prepare the command that should be written
			currentCommand = fmt.Sprintf("PX %d %d %s\n", y + *yoffset, x + *xoffset, *color)
			stripBuffer.Write([]byte(currentCommand))
		}
	}

	doneChannel <- true
	stripBufferChannel <- stripBuffer
}

// openImage opens an image at the given path and returns an image.Image
func openImage(imagePath string) image.Image {
    log.Println("[ ] Opening the image (%s) ...", imagePath)
	file, openErr := os.Open(imagePath)
    defer file.Close()
	if openErr != nil {
        log.Println("[E] Error while opening the image")
		log.Fatal(openErr)
	}
    log.Println("[ ] Done opening the image")

    log.Println("[ ] Decoding the image...")
    img, _, err := image.Decode(bufio.NewReader(file))
    if err != nil {
        log.Println("[E] Error while decoding the image")
		log.Fatal(err)
    }
    log.Println("[ ] Done decoding the image")

	return img
}

// buildSendStringImage builds the string represesenting the given image for the pixelflut server
func buildSendStringImage(image image.Image, start int, end int, doneChannel chan bool, stripBufferChannel chan bytes.Buffer) {

	var imageHeight = image.Bounds().Max.Y
	var stripBuffer bytes.Buffer

	log.Printf("ImageHeight: %d", imageHeight)

	for x := start; x < end; x++ {
		for y := 0; y < imageHeight; y++ {
			r, g, b, _ := image.At(x, y).RGBA()
            if r > 10 || g > 10 || b > 10 {
                command := fmt.Sprintf("PX %d %d 000000", x, y)
                stripBuffer.Write([]byte(command))
            } else {
                command := fmt.Sprintf("PX %d %d ffffff", x, y)
                stripBuffer.Write([]byte(command))
            }
		}
	}

	doneChannel <- true
	stripBufferChannel <- stripBuffer
}

func main() {
	log.Printf("[ ] %s -> %s:%s at (%d, %d)", *imagePath, *host, *port, xoffset, yoffset)
	parseFlags()
	testConnection()

	if *fill == true {
        log.Println("[M] FILL MODE")
		// Define channels used to bundle the data generated and get information
		doneChannel := make(chan bool)
		stripBufferChannel := make(chan bytes.Buffer)
		var completeBuffer bytes.Buffer

		// Calculate the width of the individual stripes
		var stripwidth int = *canvasWidth / *cores

		// Create a new buildStringBot generating the stripes
        count := *cores
        bar := pb.StartNew(count)
		for thread := 0; thread < *cores; thread++ {
            bar.Increment()
			go buildSendString(thread * stripwidth, (thread + 1) * stripwidth, doneChannel, stripBufferChannel)
		}
        bar.FinishPrint("Done Starting all threads")

		// Catch all threads
        count = *cores
        bar = pb.StartNew(count)
		for thread := 0; thread < *cores; thread++ {
			// get a "Done" message from each worker
			_ = <- doneChannel

			// get the buffer generated and append it to the complete buffer by writing it there
			stripBufferChannelOutput := <- stripBufferChannel
			completeBuffer.Write(stripBufferChannelOutput.Bytes())

            bar.Increment()
		}
        bar.FinishPrint("Done catching all threads")

		// Connect to the server
		connection, netDialError := net.Dial("tcp", address)
		if netDialError != nil {
			log.Fatal(netDialError)
		}

        // Write the buffer to the connection
		i := 0
        count = *loops
        bar = pb.StartNew(count)
		for i < *loops {
			// actual write
			_, writeErr := connection.Write(completeBuffer.Bytes())
			if writeErr != nil {
                log.Println("Some kind of error occured (probaly the proken pipe thing...)")
                time.Sleep(500 * time.Millisecond)
			//	log.Fatal(writeErr)
			}

            bar.Increment()
			i++
		}
        bar.FinishPrint("Done writing !")

		fmt.Printf("\n")

		// Close the connection
		connectionCloseError := connection.Close()
		if connectionCloseError != nil {
			log.Fatal(connectionCloseError)
		}

		// Cleanup
		fmt.Printf("cleanup: %d -> %d", *cores * stripwidth, *canvasWidth)
	}

	if *imagePath != "" {
        log.Println("[M] IMAGE MODE")

		// Open the image
        log.Println("[ ] Opening the image...")
		var image = openImage(*imagePath)
        log.Println("[ ] Done opening the image")

		// Define channels used to bundle the data generated and get informations
        log.Println("[ ] Creating channels...")
		doneChannel := make(chan bool)
		stripBufferChannel := make(chan bytes.Buffer)
		var completeBuffer bytes.Buffer
        log.Println("[ ] Done creating channels")

		// Calculate the width of the individual stripes
        log.Println("[ ] Creating stripwidth...")
		var stripwidth int = image.Bounds().Max.X / *cores
        log.Println("[ ] Done calculating stipwith")

		// Create new buildSendString workers building the string that should be sent to the pixelflut server
        log.Println("[ ] Creating buildworkers...")
		for thread := 0; thread < *cores; thread++ {
			log.Printf("Staring thread %d", thread)
			go buildSendStringImage(image, thread * stripwidth, (thread+1) * stripwidth, doneChannel, stripBufferChannel)
		}

		// Catch all the threads
        count := *cores
        bar := pb.StartNew(count)
		for thread := 0; thread < *cores; thread++ {
			// Get a "Done" message from each worker
			_ = <- doneChannel

			// Get the buffer generated and append it to the complete buffer by writing it there
			stripBufferChannelOutput := <- stripBufferChannel
			completeBuffer.Write(stripBufferChannelOutput.Bytes())

            bar.Increment()
		}
        bar.FinishPrint("Done catching the threads!")


		// Write the command to the server
        log.Println("[ ] Connecting to server...")
		connection, netDialError := net.Dial("tcp", address)
		if netDialError != nil {
			log.Fatal(netDialError)
		}
        log.Println("[ ] Done connecting to server")

        // Actual write
        log.Println("[ ] Actual send...")
		i := 0
        count = *loops
        bar = pb.StartNew(count)
		for i < *loops {
			_, writeErr := connection.Write(completeBuffer.Bytes())
			if writeErr != nil {
				log.Fatal(writeErr)
			}

            bar.Increment()
			i++
		}
        bar.FinishPrint("The End!")
        log.Println("[ ] Done sending")

		// Close the connection
        log.Println("[ ] Closing the connection...")
		connectionCloseError := connection.Close()
		if connectionCloseError != nil {
			log.Fatal(connectionCloseError)
		}
        log.Println("[ ] Done closing the connection")

		// Cleanup
		log.Printf("[ ] Cleanup: %d -> %d", *cores * stripwidth, *canvasWidth)
	}
}