package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"flag"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"os"
)

// OptimizedIdenticon with indexed colors for smaller file sizes
type OptimizedIdenticon struct {
	source []byte
	size   int
}

// NewOptimizedIdenticon creates a generator with indexed colors
func NewOptimizedIdenticon(source []byte) *OptimizedIdenticon {
	return &OptimizedIdenticon{
		source: source,
		size:   256,
	}
}

// NewOptimizedIdenticonWithSize creates a generator with custom size
func NewOptimizedIdenticonWithSize(source []byte, size int) *OptimizedIdenticon {
	return &OptimizedIdenticon{
		source: source,
		size:   size,
	}
}

// getBit returns the n-th bit (0-indexed) from source
func (identicon *OptimizedIdenticon) getBit(n int) bool {
	if len(identicon.source) == 0 || n < 0 {
		return false
	}
	byteIndex := n / 8
	bitIndex := n % 8
	if byteIndex >= len(identicon.source) {
		return false
	}
	return (identicon.source[byteIndex]>>bitIndex)&1 == 1
}

// getByte returns the n-th byte, wraps around if needed
func (identicon *OptimizedIdenticon) getByte(n int) byte {
	if len(identicon.source) == 0 {
		return 0
	}
	return identicon.source[n%len(identicon.source)]
}

// getColorIndices returns color indices for indexed version
func (identicon *OptimizedIdenticon) getColorIndices() (primaryIndex, secondaryIndex, bgIndex uint8) {
	if len(identicon.source) < 32 {
		return 0, 1, 2
	}

	// Primary color index (4 bits → 16 colors)
	primaryIndex = 0
	for i := 0; i < 4; i++ {
		if identicon.getBit(248 + i) {
			primaryIndex |= 1 << i
		}
	}
	primaryIndex %= 16

	// Secondary color index (4 bits → 16 colors)
	secondaryIndex = 0
	for i := 0; i < 4; i++ {
		if identicon.getBit(244 + i) {
			secondaryIndex |= 1 << i
		}
	}
	secondaryIndex %= 16

	// Background choice (2 bits → 4 options)
	bgChoice := 0
	for i := 0; i < 2; i++ {
		if identicon.getBit(252 + i) {
			bgChoice |= 1 << i
		}
	}
	bgIndex = uint8(bgChoice % 3) // 0, 1, or 2

	return primaryIndex, secondaryIndex, bgIndex
}

// generatePixelPattern generates 5x5 symmetric pixel grid
func (identicon *OptimizedIdenticon) generatePixelPattern() ([]bool, []bool) {
	primary := make([]bool, 25)
	secondary := make([]bool, 25)

	// Use bits 0-14 for primary pattern
	bitIndex := 0
	for row := 0; row < 5; row++ {
		for col := 0; col < 3; col++ {
			paint := identicon.getBit(bitIndex)
			bitIndex++

			ix := row*5 + col
			mirrorIx := row*5 + (4 - col)
			primary[ix] = paint
			primary[mirrorIx] = paint
		}
	}

	// Use bits 15-29 for secondary pattern
	for row := 0; row < 5; row++ {
		for col := 0; col < 3; col++ {
			paint := identicon.getBit(bitIndex)
			bitIndex++

			ix := row*5 + col
			mirrorIx := row*5 + (4 - col)
			secondary[ix] = paint
			secondary[mirrorIx] = paint
		}
	}

	return primary, secondary
}

// createPalette creates an optimized palette with only necessary colors
func createPalette(primaryIdx, secondaryIdx, bgIdx uint8, darkMode bool) color.Palette {
	// Optimized color palettes - only 16 colors per palette
	primaryPalette := []color.Color{
		color.RGBA{0x00, 0xbf, 0x93, 0xff}, // turquoise
		color.RGBA{0x2d, 0xcc, 0x70, 0xff}, // mint
		color.RGBA{0x42, 0xe4, 0x53, 0xff}, // green
		color.RGBA{0xf1, 0xc4, 0x0f, 0xff}, // yellowOrange
		color.RGBA{0xe6, 0x7f, 0x22, 0xff}, // brown
		color.RGBA{0xff, 0x94, 0x4e, 0xff}, // orange
		color.RGBA{0xe8, 0x4c, 0x3d, 0xff}, // red
		color.RGBA{0x35, 0x98, 0xdb, 0xff}, // blue
		color.RGBA{0x9a, 0x59, 0xb5, 0xff}, // purple
		color.RGBA{0xef, 0x3e, 0x96, 0xff}, // magenta
		color.RGBA{0xdf, 0x21, 0xb9, 0xff}, // violet
		color.RGBA{0x7d, 0xc2, 0xd2, 0xff}, // lightBlue
		color.RGBA{0x16, 0xa0, 0x86, 0xff}, // turquoiseIntense
		color.RGBA{0x27, 0xae, 0x61, 0xff}, // mintIntense
		color.RGBA{0x24, 0xc3, 0x33, 0xff}, // greenIntense
		color.RGBA{0x1c, 0xab, 0xbb, 0xff}, // lightBlueIntense
	}

	secondaryPalette := []color.Color{
		color.RGBA{0x34, 0x49, 0x5e, 0xff}, // darkBlue
		color.RGBA{0x95, 0xa5, 0xa5, 0xff}, // grey
		color.RGBA{0xd2, 0x54, 0x00, 0xff}, // brownIntense
		color.RGBA{0xc1, 0x39, 0x2b, 0xff}, // redIntense
		color.RGBA{0x29, 0x7f, 0xb8, 0xff}, // blueIntense
		color.RGBA{0x8d, 0x44, 0xad, 0xff}, // purpleIntense
		color.RGBA{0xbe, 0x12, 0x7e, 0xff}, // violetIntense
		color.RGBA{0xe5, 0x23, 0x83, 0xff}, // magentaIntense
		color.RGBA{0x27, 0xae, 0x61, 0xff}, // mintIntense
		color.RGBA{0x24, 0xc3, 0x33, 0xff}, // greenIntense
		color.RGBA{0xd9, 0xd9, 0x21, 0xff}, // yellowIntense
		color.RGBA{0xf3, 0x9c, 0x11, 0xff}, // yellowOrangeIntense
		color.RGBA{0xff, 0x55, 0x00, 0xff}, // orangeIntense
		color.RGBA{0x1c, 0xab, 0xbb, 0xff}, // lightBlueIntense
		color.RGBA{0x23, 0x23, 0x23, 0xff}, // lightBlackIntense
		color.RGBA{0x7e, 0x8c, 0x8d, 0xff}, // greyIntense
	}

	// Background colors based on mode
	lightBackgrounds := []color.Color{
		color.RGBA{255, 255, 255, 255}, // white
		color.RGBA{243, 245, 247, 255}, // light gray 1
		color.RGBA{236, 240, 241, 255}, // light gray 2
		color.RGBA{0, 0, 0, 0},         // transparent (position 3)
	}

	darkBackgrounds := []color.Color{
		color.RGBA{30, 30, 30, 255},  // dark gray
		color.RGBA{45, 62, 80, 255},  // dark blue
		color.RGBA{57, 57, 57, 255},  // dark gray 2
		color.RGBA{0, 0, 0, 0},       // transparent (position 3)
	}

	// Palette in exact order:
	// 0: Background
	// 1: Primary color
	// 2: Secondary color
	// 3: Transparent (optional)
	palette := make(color.Palette, 0, 4)

	// Background first
	if darkMode {
		palette = append(palette, darkBackgrounds[bgIdx])
	} else {
		palette = append(palette, lightBackgrounds[bgIdx])
	}

	// Then primary and secondary colors
	palette = append(palette,
		primaryPalette[primaryIdx],
		secondaryPalette[secondaryIdx],
		color.RGBA{0, 0, 0, 0}, // transparent as last option
	)

	return palette
}

// Generate48x48ForFace creates a 48x48 pixel image specifically for Face headers
func (identicon *OptimizedIdenticon) Generate48x48ForFace(transparent bool) *image.Paletted {
	const (
		size       = 48
		spriteSize = 5
		pixelSize  = 6                             // 48/8 = 6
		margin     = (size - pixelSize*spriteSize) / 2 // = 9
	)

	// Determine color indices
	primaryIdx, secondaryIdx, bgIdx := identicon.getColorIndices()

	// For transparent background, set bgIdx to 3 (transparent)
	if transparent {
		bgIdx = 3
	}

	// Palette for export (always light mode for better compatibility)
	palette := createPalette(primaryIdx, secondaryIdx, bgIdx, false)

	// Create image
	img := image.NewPaletted(image.Rect(0, 0, size, size), palette)

	// Fill background
	bgIndex := uint8(0)
	if transparent {
		bgIndex = 3 // transparent
	}

	for i := 0; i < size; i++ {
		for j := 0; j < size; j++ {
			img.SetColorIndex(j, i, bgIndex)
		}
	}

	primaryPixels, secondaryPixels := identicon.generatePixelPattern()

	// Secondary pixels (index 2)
	for row := 0; row < spriteSize; row++ {
		for col := 0; col < spriteSize; col++ {
			if secondaryPixels[row*spriteSize+col] {
				x := col*pixelSize + margin
				y := row*pixelSize + margin
				for py := y; py < y+pixelSize; py++ {
					for px := x; px < x+pixelSize; px++ {
						if px < size && py < size {
							img.SetColorIndex(px, py, 2)
						}
					}
				}
			}
		}
	}

	// Primary pixels (index 1)
	for row := 0; row < spriteSize; row++ {
		for col := 0; col < spriteSize; col++ {
			if primaryPixels[row*spriteSize+col] {
				x := col*pixelSize + margin
				y := row*pixelSize + margin
				for py := y; py < y+pixelSize; py++ {
					for px := x; px < x+pixelSize; px++ {
						if px < size && py < size {
							img.SetColorIndex(px, py, 1)
						}
					}
				}
			}
		}
	}

	return img
}

// GenerateForExportOptimized for indexed export (256x256)
func (identicon *OptimizedIdenticon) GenerateForExportOptimized(transparent bool) *image.Paletted {
	const (
		spriteSize = 5
	)

	pixelSize := identicon.size / 8
	margin := (identicon.size - pixelSize*spriteSize) / 2

	// Determine color indices
	primaryIdx, secondaryIdx, bgIdx := identicon.getColorIndices()

	// For transparent background, set bgIdx to 3 (transparent)
	if transparent {
		bgIdx = 3
	}

	// Palette for export (always light mode for better compatibility)
	palette := createPalette(primaryIdx, secondaryIdx, bgIdx, false)

	// Create image
	img := image.NewPaletted(image.Rect(0, 0, identicon.size, identicon.size), palette)

	// Fill background
	bgIndex := uint8(0)
	if transparent {
		bgIndex = 3 // transparent
	}

	for i := 0; i < identicon.size; i++ {
		for j := 0; j < identicon.size; j++ {
			img.SetColorIndex(j, i, bgIndex)
		}
	}

	primaryPixels, secondaryPixels := identicon.generatePixelPattern()

	// Secondary pixels (index 2)
	for row := 0; row < spriteSize; row++ {
		for col := 0; col < spriteSize; col++ {
			if secondaryPixels[row*spriteSize+col] {
				x := col*pixelSize + margin
				y := row*pixelSize + margin
				for py := y; py < y+pixelSize; py++ {
					for px := x; px < x+pixelSize; px++ {
						if px < identicon.size && py < identicon.size {
							img.SetColorIndex(px, py, 2)
						}
					}
				}
			}
		}
	}

	// Primary pixels (index 1)
	for row := 0; row < spriteSize; row++ {
		for col := 0; col < spriteSize; col++ {
			if primaryPixels[row*spriteSize+col] {
				x := col*pixelSize + margin
				y := row*pixelSize + margin
				for py := y; py < y+pixelSize; py++ {
					for px := x; px < x+pixelSize; px++ {
						if px < identicon.size && py < identicon.size {
							img.SetColorIndex(px, py, 1)
						}
					}
				}
			}
		}
	}

	return img
}

func main() {
	// CLI flags
	input := flag.String("input", "", "Input text (username|email|pubkey)")
	size := flag.Int("size", 48, "Image size (48 or 256)")
	transparent := flag.Bool("transparent", true, "Transparent background")
	outputFormat := flag.String("format", "base64", "Output format: base64, dataurl, or png")
	outputFile := flag.String("output", "", "Output file (for png format)")
	hash := flag.String("hash", "", "Direct SHA256 hash (hex) instead of input")

	flag.Parse()

	// Validate input
	if *input == "" && *hash == "" {
		fmt.Fprintln(os.Stderr, "Error: -input or -hash required")
		fmt.Fprintln(os.Stderr, "Usage: identicons-cli -input 'username|email|pubkey' [-size 48|256] [-transparent] [-format base64|dataurl|png] [-output file.png]")
		os.Exit(1)
	}

	// Generate hash
	var hashBytes []byte
	if *hash != "" {
		// Use provided hash (hex string)
		fmt.Sscanf(*hash, "%x", &hashBytes)
		if len(hashBytes) != 32 {
			fmt.Fprintln(os.Stderr, "Error: hash must be 32 bytes (64 hex chars)")
			os.Exit(1)
		}
	} else {
		// Hash the input
		h := sha256.Sum256([]byte(*input))
		hashBytes = h[:]
	}

	// Generate identicon
	var img *image.Paletted
	if *size == 48 {
		identicon := NewOptimizedIdenticonWithSize(hashBytes, 48)
		img = identicon.Generate48x48ForFace(*transparent)
	} else if *size == 256 {
		identicon := NewOptimizedIdenticonWithSize(hashBytes, 256)
		img = identicon.GenerateForExportOptimized(*transparent)
	} else {
		fmt.Fprintln(os.Stderr, "Error: size must be 48 or 256")
		os.Exit(1)
	}

	// Encode to PNG
	var buf bytes.Buffer
	encoder := png.Encoder{
		CompressionLevel: png.BestCompression,
	}

	if err := encoder.Encode(&buf, img); err != nil {
		fmt.Fprintln(os.Stderr, "Error encoding PNG:", err)
		os.Exit(1)
	}

	// Output based on format
	switch *outputFormat {
	case "base64":
		// Output base64 only (for Face header)
		b64 := base64.StdEncoding.EncodeToString(buf.Bytes())
		fmt.Print(b64)

	case "dataurl":
		// Output data URL (for web)
		b64 := base64.StdEncoding.EncodeToString(buf.Bytes())
		fmt.Printf("data:image/png;base64,%s", b64)

	case "png":
		// Output PNG file
		if *outputFile == "" {
			fmt.Fprintln(os.Stderr, "Error: -output required for png format")
			os.Exit(1)
		}
		if err := os.WriteFile(*outputFile, buf.Bytes(), 0644); err != nil {
			fmt.Fprintln(os.Stderr, "Error writing PNG:", err)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stderr, "PNG saved to %s (%d bytes)\n", *outputFile, buf.Len())

	default:
		fmt.Fprintln(os.Stderr, "Error: format must be base64, dataurl, or png")
		os.Exit(1)
	}
}