// rim.go
// read as "r img o" as in "random image output", "rimgo"
package main

import (
	"image"
	"image/color"
	"image/png"
	"log"
	"math/rand"
	"os"
	"time"

    // my depends
    "bytes"
    "golang.org/x/term"
    "github.com/schollz/progressbar/v3"
    go_col "github.com/TwiN/go-color"
    "github.com/klauspost/compress/gzip"
)

var width  int = 10000
var height int = 10000

func Gen(n int) uint8 {
	return uint8(rand.Intn(n))
}

func handleError(err error, ux_error string) {
    if term.IsTerminal(int(os.Stdout.Fd())) {
        println(go_col.InRed(ux_error))
    } else {
        println(ux_error)
    }
    log.Fatal(err)
}

func imageToRGBA(img image.Image) []uint8 {
    sz := img.Bounds()
    raw := make([]uint8, (sz.Max.X-sz.Min.X)*(sz.Max.Y-sz.Min.Y)*4)
    idx := 0
    for y := sz.Min.Y; y < sz.Max.Y; y++ {
        for x := sz.Min.X; x < sz.Max.X; x++ {
            r, g, b, a := img.At(x, y).RGBA()
            raw[idx], raw[idx+1], raw[idx+2], raw[idx+3] = uint8(r), uint8(g), uint8(b), uint8(a)
            idx += 4
        }
    }
    return raw
}

func main() {
    // added code
    bar := progressbar.Default(int64(width))

	rand.Seed(time.Now().UnixNano())

	var rb int = int(Gen(0xff))
	var gb int = int(Gen(0xff))
	var bb int = int(Gen(0xff))

	img := image.NewRGBA(image.Rect(0,0,width,height))

	for x := 0; x < width; x++ {
		for y := 0; y < height; y++ {
			img.Set(x, y, color.RGBA{Gen(rb), Gen(gb), Gen(bb), 0xff})
		}
        bar.Add(1)
	}

    println(go_col.InGreen("Image generated!"))
	f, err := os.Create("output.png")
	if err != nil {
        handleError(err, "Error creating output file")
	}

    println(go_col.InGreen("Writing image..."))
	if err := png.Encode(f, img); err != nil {
		f.Close()
        handleError(err, "Error writing output file")
	}

	if err := f.Close(); err != nil {
        handleError(err, "Error closing output file")
	}
    println(go_col.InGreen("Image saved!"))

    // Compressing raw data and writing to output.raw.gz
    // buffer for gzip
    var buf bytes.Buffer
	zw := gzip.NewWriter(&buf)
	// Setting header fields
	zw.Name = "output.png"
	zw.Comment = "Image generated by rimgo AI"
	_, err = zw.Write(imageToRGBA(img))
    zw.Close()
    println(go_col.InGreen("Saving raw gzipped colors"))
    file, err := os.Create("output.raw.gz")
    println(go_col.InGreen("Raw gzipped colors saved!"))
	if err != nil {
        handleError(err, "Error saving output.raw.gz (raw image bytes)")
	}
	defer file.Close()
	// Write the buffer data to the file
    buf.WriteTo(file)
	if err != nil {
		panic(err)
	}
    println(go_col.InGreen("Done!"))
	// Close the file
}
