package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"image"
	"image/png"
	"log"
	"math"
	"os"
	"path/filepath"
	"strings"
)

// --- Configuration ---
var DPI float64 = 1000.0 // Higher DPI = smoother curves
var PixelToMM float64 = 25.4 / DPI

var StencilHeight float64 = 0.16 // mm, default
var WallHeight float64 = 2.0     // mm, default
var WallThickness float64 = 1.0  // mm, default
var KeepPNG bool

// --- STL Helpers ---

type Point struct {
	X, Y, Z float64
}

func WriteSTL(filename string, triangles [][3]Point) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	// Write Binary STL Header (80 bytes)
	header := make([]byte, 80)
	copy(header, "Generated by pcb-to-stencil")
	if _, err := f.Write(header); err != nil {
		return err
	}

	// Write Number of Triangles (4 bytes uint32)
	count := uint32(len(triangles))
	if err := binary.Write(f, binary.LittleEndian, count); err != nil {
		return err
	}

	// Write Triangles
	// Each triangle is 50 bytes:
	// Normal (3 floats = 12 bytes)
	// Vertex 1 (3 floats = 12 bytes)
	// Vertex 2 (3 floats = 12 bytes)
	// Vertex 3 (3 floats = 12 bytes)
	// Attribute byte count (2 bytes uint16)

	// Buffer for a single triangle to minimize syscalls
	buf := make([]byte, 50)

	for _, t := range triangles {
		// Normal (0,0,0)
		binary.LittleEndian.PutUint32(buf[0:4], math.Float32bits(0))
		binary.LittleEndian.PutUint32(buf[4:8], math.Float32bits(0))
		binary.LittleEndian.PutUint32(buf[8:12], math.Float32bits(0))

		// Vertex 1
		binary.LittleEndian.PutUint32(buf[12:16], math.Float32bits(float32(t[0].X)))
		binary.LittleEndian.PutUint32(buf[16:20], math.Float32bits(float32(t[0].Y)))
		binary.LittleEndian.PutUint32(buf[20:24], math.Float32bits(float32(t[0].Z)))

		// Vertex 2
		binary.LittleEndian.PutUint32(buf[24:28], math.Float32bits(float32(t[1].X)))
		binary.LittleEndian.PutUint32(buf[28:32], math.Float32bits(float32(t[1].Y)))
		binary.LittleEndian.PutUint32(buf[32:36], math.Float32bits(float32(t[1].Z)))

		// Vertex 3
		binary.LittleEndian.PutUint32(buf[36:40], math.Float32bits(float32(t[2].X)))
		binary.LittleEndian.PutUint32(buf[40:44], math.Float32bits(float32(t[2].Y)))
		binary.LittleEndian.PutUint32(buf[44:48], math.Float32bits(float32(t[2].Z)))

		// Attribute byte count (0)
		binary.LittleEndian.PutUint16(buf[48:50], 0)

		if _, err := f.Write(buf); err != nil {
			return err
		}
	}
	return nil
}

func AddBox(triangles *[][3]Point, x, y, w, h, zHeight float64) {
	x0, y0 := x, y
	x1, y1 := x+w, y+h
	z0, z1 := 0.0, zHeight

	p000 := Point{x0, y0, z0}
	p100 := Point{x1, y0, z0}
	p110 := Point{x1, y1, z0}
	p010 := Point{x0, y1, z0}
	p001 := Point{x0, y0, z1}
	p101 := Point{x1, y0, z1}
	p111 := Point{x1, y1, z1}
	p011 := Point{x0, y1, z1}

	addQuad := func(a, b, c, d Point) {
		*triangles = append(*triangles, [3]Point{a, b, c})
		*triangles = append(*triangles, [3]Point{c, d, a})
	}

	addQuad(p000, p010, p110, p100) // Bottom
	addQuad(p101, p111, p011, p001) // Top
	addQuad(p000, p100, p101, p001) // Front
	addQuad(p100, p110, p111, p101) // Right
	addQuad(p110, p010, p011, p111) // Back
	addQuad(p010, p000, p001, p011) // Left
}

// --- Meshing Logic (Optimized) ---

// ComputeWallMask generates a mask for the wall based on the outline image.
// It identifies the board area (inside the outline) and creates a wall of
// specified thickness around it.
func ComputeWallMask(img image.Image, thicknessMM float64) ([]bool, []bool) {
	bounds := img.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	size := w * h

	// Helper for neighbors
	dx := []int{0, 0, 1, -1}
	dy := []int{1, -1, 0, 0}

	// 1. Identify Outline Pixels (White)
	isOutline := make([]bool, size)
	outlineQueue := []int{}
	for i := 0; i < size; i++ {
		cx := i % w
		cy := i / w
		c := img.At(cx, cy)
		r, _, _, _ := c.RGBA()
		if r > 10000 { // White-ish
			isOutline[i] = true
			outlineQueue = append(outlineQueue, i)
		}
	}

	// 2. Dilate Outline to close gaps
	// We dilate by a small amount (e.g. 0.5mm) to ensure the outline is closed.
	gapClosingMM := 0.5
	gapClosingPixels := int(gapClosingMM / PixelToMM)
	if gapClosingPixels < 1 {
		gapClosingPixels = 1
	}

	dist := make([]int, size)
	for i := 0; i < size; i++ {
		if isOutline[i] {
			dist[i] = 0
		} else {
			dist[i] = -1
		}
	}

	// BFS for Dilation
	dilatedOutline := make([]bool, size)
	copy(dilatedOutline, isOutline)

	// Use a separate queue for dilation to avoid modifying the original outlineQueue if we needed it
	dQueue := make([]int, len(outlineQueue))
	copy(dQueue, outlineQueue)

	for len(dQueue) > 0 {
		idx := dQueue[0]
		dQueue = dQueue[1:]

		d := dist[idx]
		if d >= gapClosingPixels {
			continue
		}

		cx := idx % w
		cy := idx / w

		for i := 0; i < 4; i++ {
			nx, ny := cx+dx[i], cy+dy[i]
			if nx >= 0 && nx < w && ny >= 0 && ny < h {
				nIdx := ny*w + nx
				if dist[nIdx] == -1 {
					dist[nIdx] = d + 1
					dilatedOutline[nIdx] = true
					dQueue = append(dQueue, nIdx)
				}
			}
		}
	}

	// 3. Flood Fill "Outside" using Dilated Outline as barrier
	isOutside := make([]bool, size)
	// Start from (0,0) - assumed to be outside due to padding
	if !dilatedOutline[0] {
		isOutside[0] = true
		fQueue := []int{0}

		for len(fQueue) > 0 {
			idx := fQueue[0]
			fQueue = fQueue[1:]

			cx := idx % w
			cy := idx / w

			for i := 0; i < 4; i++ {
				nx, ny := cx+dx[i], cy+dy[i]
				if nx >= 0 && nx < w && ny >= 0 && ny < h {
					nIdx := ny*w + nx
					if !isOutside[nIdx] && !dilatedOutline[nIdx] {
						isOutside[nIdx] = true
						fQueue = append(fQueue, nIdx)
					}
				}
			}
		}
	}

	// 4. Restore Board Shape (Erode "Outside" back to original boundary)
	// We dilated the outline, so "Outside" stopped 'gapClosingPixels' away from the real board edge.
	// We need to expand "Outside" inwards by 'gapClosingPixels' to touch the real board edge.
	// Then "Board" = !Outside.

	// Reset dist for Outside expansion
	for i := 0; i < size; i++ {
		if isOutside[i] {
			dist[i] = 0
		} else {
			dist[i] = -1
		}
	}

	oQueue := []int{}
	for i := 0; i < size; i++ {
		if isOutside[i] {
			oQueue = append(oQueue, i)
		}
	}

	isOutsideExpanded := make([]bool, size)
	copy(isOutsideExpanded, isOutside)

	for len(oQueue) > 0 {
		idx := oQueue[0]
		oQueue = oQueue[1:]

		d := dist[idx]
		if d >= gapClosingPixels {
			continue
		}

		cx := idx % w
		cy := idx / w

		for i := 0; i < 4; i++ {
			nx, ny := cx+dx[i], cy+dy[i]
			if nx >= 0 && nx < w && ny >= 0 && ny < h {
				nIdx := ny*w + nx
				if dist[nIdx] == -1 {
					dist[nIdx] = d + 1
					isOutsideExpanded[nIdx] = true
					oQueue = append(oQueue, nIdx)
				}
			}
		}
	}

	// 5. Define Board
	isBoard := make([]bool, size)
	for i := 0; i < size; i++ {
		isBoard[i] = !isOutsideExpanded[i]
	}

	// 6. Generate Wall
	// Wall is generated by expanding Board outwards.
	// We want the wall to be strictly OUTSIDE the board (or centered on outline? User said "starts at outline").
	// If we expand Board, we get pixels outside.

	thicknessPixels := int(thicknessMM / PixelToMM)
	if thicknessPixels < 1 {
		thicknessPixels = 1
	}

	// Reset dist for Wall generation
	for i := 0; i < size; i++ {
		if isBoard[i] {
			dist[i] = 0
		} else {
			dist[i] = -1
		}
	}

	wQueue := []int{}
	for i := 0; i < size; i++ {
		if isBoard[i] {
			wQueue = append(wQueue, i)
		}
	}

	isWall := make([]bool, size)

	for len(wQueue) > 0 {
		idx := wQueue[0]
		wQueue = wQueue[1:]

		d := dist[idx]
		if d >= thicknessPixels {
			continue
		}

		cx := idx % w
		cy := idx / w

		for i := 0; i < 4; i++ {
			nx, ny := cx+dx[i], cy+dy[i]
			if nx >= 0 && nx < w && ny >= 0 && ny < h {
				nIdx := ny*w + nx
				if dist[nIdx] == -1 {
					dist[nIdx] = d + 1
					isWall[nIdx] = true
					wQueue = append(wQueue, nIdx)
				}
			}
		}
	}

	return isWall, isBoard
}

func GenerateMeshFromImages(stencilImg, outlineImg image.Image) [][3]Point {
	bounds := stencilImg.Bounds()
	width := bounds.Max.X
	height := bounds.Max.Y
	var triangles [][3]Point

	var wallMask []bool
	var boardMask []bool
	if outlineImg != nil {
		fmt.Println("Computing wall mask...")
		wallMask, boardMask = ComputeWallMask(outlineImg, WallThickness)
	}

	// Optimization: Run-Length Encoding
	for y := 0; y < height; y++ {
		var startX = -1
		var currentHeight = 0.0

		for x := 0; x < width; x++ {
			// Check stencil (black = solid)
			sc := stencilImg.At(x, y)
			sr, sg, sb, _ := sc.RGBA()
			isStencilSolid := sr < 10000 && sg < 10000 && sb < 10000

			// Check wall
			isWall := false
			isInsideBoard := true
			if wallMask != nil {
				idx := y*width + x
				isWall = wallMask[idx]
				if boardMask != nil {
					isInsideBoard = boardMask[idx]
				}
			}

			// Determine height at this pixel
			h := 0.0
			if isWall {
				h = WallHeight
			} else if isStencilSolid {
				if isInsideBoard {
					h = StencilHeight
				}
			}

			if h > 0 {
				if startX == -1 {
					startX = x
					currentHeight = h
				} else if h != currentHeight {
					// Height changed, end current strip and start new one
					stripLen := x - startX
					AddBox(
						&triangles,
						float64(startX)*PixelToMM,
						float64(y)*PixelToMM,
						float64(stripLen)*PixelToMM,
						PixelToMM,
						currentHeight,
					)
					startX = x
					currentHeight = h
				}
			} else {
				if startX != -1 {
					// End of strip, generate box
					stripLen := x - startX
					AddBox(
						&triangles,
						float64(startX)*PixelToMM,
						float64(y)*PixelToMM,
						float64(stripLen)*PixelToMM,
						PixelToMM,
						currentHeight,
					)
					startX = -1
					currentHeight = 0.0
				}
			}
		}
		if startX != -1 {
			stripLen := width - startX
			AddBox(
				&triangles,
				float64(startX)*PixelToMM,
				float64(y)*PixelToMM,
				float64(stripLen)*PixelToMM,
				PixelToMM,
				currentHeight,
			)
		}
	}
	return triangles
}

// --- Main ---

func main() {
	flag.Float64Var(&StencilHeight, "height", 0.16, "Stencil height in mm")
	flag.Float64Var(&WallHeight, "wall-height", 2.0, "Wall height in mm")
	flag.Float64Var(&WallThickness, "wall-thickness", 1, "Wall thickness in mm")
	flag.Float64Var(&DPI, "dpi", 1000.0, "DPI for rendering (lower = smaller file, rougher curves)")
	flag.BoolVar(&KeepPNG, "keep-png", false, "Save intermediate PNG file")
	flag.Parse()

	// Update PixelToMM based on DPI flag
	PixelToMM = 25.4 / DPI

	args := flag.Args()
	if len(args) < 1 {
		fmt.Println("Usage: go run main.go [options] <path_to_gerber_file> [path_to_outline_gerber_file]")
		fmt.Println("Options:")
		flag.PrintDefaults()
		fmt.Println("Example: go run main.go -height=0.3 MyPCB.GTP MyPCB.GKO")
		os.Exit(1)
	}

	gerberPath := args[0]
	var outlinePath string
	if len(args) > 1 {
		outlinePath = args[1]
	}

	outputPath := strings.TrimSuffix(gerberPath, filepath.Ext(gerberPath)) + ".stl"

	// 1. Parse Gerber(s)
	fmt.Printf("Parsing %s...\n", gerberPath)
	gf, err := ParseGerber(gerberPath)
	if err != nil {
		log.Fatalf("Error parsing gerber: %v", err)
	}

	var outlineGf *GerberFile
	if outlinePath != "" {
		fmt.Printf("Parsing outline %s...\n", outlinePath)
		outlineGf, err = ParseGerber(outlinePath)
		if err != nil {
			log.Fatalf("Error parsing outline gerber: %v", err)
		}
	}

	// 2. Calculate Union Bounds
	bounds := gf.CalculateBounds()
	if outlineGf != nil {
		outlineBounds := outlineGf.CalculateBounds()
		if outlineBounds.MinX < bounds.MinX {
			bounds.MinX = outlineBounds.MinX
		}
		if outlineBounds.MinY < bounds.MinY {
			bounds.MinY = outlineBounds.MinY
		}
		if outlineBounds.MaxX > bounds.MaxX {
			bounds.MaxX = outlineBounds.MaxX
		}
		if outlineBounds.MaxY > bounds.MaxY {
			bounds.MaxY = outlineBounds.MaxY
		}
	}

	// Expand bounds to accommodate wall thickness and prevent clipping
	// We add WallThickness + extra margin to all sides
	margin := WallThickness + 5.0 // mm
	bounds.MinX -= margin
	bounds.MinY -= margin
	bounds.MaxX += margin
	bounds.MaxY += margin

	// 3. Render to Image(s)
	fmt.Println("Rendering to internal image...")
	img := gf.Render(DPI, &bounds)

	var outlineImg image.Image
	if outlineGf != nil {
		fmt.Println("Rendering outline to internal image...")
		outlineImg = outlineGf.Render(DPI, &bounds)
	}

	if KeepPNG {
		pngPath := strings.TrimSuffix(gerberPath, filepath.Ext(gerberPath)) + ".png"
		fmt.Printf("Saving intermediate PNG to %s...\n", pngPath)
		f, err := os.Create(pngPath)
		if err != nil {
			log.Printf("Warning: Could not create PNG file: %v", err)
		} else {
			if err := png.Encode(f, img); err != nil {
				log.Printf("Warning: Could not encode PNG: %v", err)
			}
			f.Close()
		}

		if outlineImg != nil {
			outlinePngPath := strings.TrimSuffix(gerberPath, filepath.Ext(gerberPath)) + "_outline.png"
			fmt.Printf("Saving intermediate Outline PNG to %s...\n", outlinePngPath)
			f, err := os.Create(outlinePngPath)
			if err != nil {
				log.Printf("Warning: Could not create Outline PNG file: %v", err)
			} else {
				if err := png.Encode(f, outlineImg); err != nil {
					log.Printf("Warning: Could not encode Outline PNG: %v", err)
				}
				f.Close()
			}
		}
	}

	// 4. Generate Mesh
	fmt.Println("Generating mesh (this may take 10-20 seconds for large boards)...")
	triangles := GenerateMeshFromImages(img, outlineImg)

	// 5. Save STL
	fmt.Printf("Saving to %s (%d triangles)...\n", outputPath, len(triangles))
	err = WriteSTL(outputPath, triangles)
	if err != nil {
		log.Fatalf("Error writing STL: %v", err)
	}

	fmt.Println("Success! Happy printing.")
}