package main

import (
	"fmt"
	"image"
	"image/png"
	"log"
	"os"
	"path/filepath"
	"strings"
)

// Config holds stencil generation parameters
type Config struct {
	StencilHeight float64
	WallHeight    float64
	WallThickness float64
	LineWidth     float64
	DPI           float64
	KeepPNG       bool
}

// Default values
const (
	DefaultStencilHeight = 0.16
	DefaultWallHeight    = 2.0
	DefaultWallThickness = 1.0
	DefaultDPI           = 1000.0
)

// ComputeWallMask generates a mask for the wall based on the outline image.
func ComputeWallMask(img image.Image, thicknessMM float64, pixelToMM float64) ([]int, []bool) {
	bounds := img.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	size := w * h

	dx := []int{0, 0, 1, -1}
	dy := []int{1, -1, 0, 0}

	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 {
			isOutline[i] = true
			outlineQueue = append(outlineQueue, i)
		}
	}

	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
		}
	}

	dilatedOutline := make([]bool, size)
	copy(dilatedOutline, isOutline)

	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)
				}
			}
		}
	}

	isOutside := make([]bool, size)
	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)
					}
				}
			}
		}
	}

	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)
				}
			}
		}
	}

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

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

	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)
		}
	}

	wallDist := make([]int, size)
	for i := range wallDist {
		wallDist[i] = -1
	}

	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
					wallDist[nIdx] = d + 1
					wQueue = append(wQueue, nIdx)
				}
			}
		}
	}

	return wallDist, isBoard
}

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

	var wallDist []int
	var boardMask []bool
	if outlineImg != nil {
		wallDist, boardMask = ComputeWallMask(outlineImg, cfg.WallThickness, pixelToMM)
	}

	for y := 0; y < height; y++ {
		var startX = -1
		var currentHeight = 0.0

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

			isWall := false
			isInsideBoard := true
			if wallDist != nil {
				idx := y*width + x
				isWall = wallDist[idx] >= 0
				if boardMask != nil {
					isInsideBoard = boardMask[idx]
				}
			}

			h := 0.0
			if isWall {
				h = cfg.WallHeight
			} else if isStencilSolid {
				if isInsideBoard {
					h = cfg.WallHeight
				}
			}

			if h > 0 {
				if startX == -1 {
					startX = x
					currentHeight = h
				} else if h != currentHeight {
					stripLen := x - startX
					AddBox(&triangles, float64(startX)*pixelToMM, float64(y)*pixelToMM,
						float64(stripLen)*pixelToMM, pixelToMM, currentHeight)
					startX = x
					currentHeight = h
				}
			} else {
				if startX != -1 {
					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
}

// processPCB handles stencil generation from gerber files
func processPCB(gerberPath, outlinePath string, cfg Config, exports []string) ([]string, image.Image, image.Image, error) {
	baseName := strings.TrimSuffix(gerberPath, filepath.Ext(gerberPath))
	var generatedFiles []string

	wantsType := func(t string) bool {
		for _, e := range exports {
			if e == t {
				return true
			}
		}
		return false
	}

	if len(exports) == 0 {
		exports = []string{"stl"}
	}

	fmt.Printf("Parsing %s...\n", gerberPath)
	gf, err := ParseGerber(gerberPath)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("error parsing gerber: %v", err)
	}

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

	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
		}
	}

	margin := cfg.WallThickness + 5.0
	bounds.MinX -= margin
	bounds.MinY -= margin
	bounds.MaxX += margin
	bounds.MaxY += margin

	fmt.Println("Rendering to internal image...")
	img := gf.Render(cfg.DPI, &bounds)

	var outlineImg image.Image
	if outlineGf != nil {
		outlineImg = outlineGf.Render(cfg.DPI, &bounds)
	}

	if cfg.KeepPNG {
		pngPath := strings.TrimSuffix(gerberPath, filepath.Ext(gerberPath)) + ".png"
		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()
		}
	}

	var triangles [][3]Point
	if wantsType("stl") {
		fmt.Println("Generating mesh...")
		triangles = GenerateMeshFromImages(img, outlineImg, cfg)
	}

	if wantsType("stl") {
		outputFilename := baseName + ".stl"
		fmt.Printf("Saving to %s (%d triangles)...\n", outputFilename, len(triangles))
		if err := WriteSTL(outputFilename, triangles); err != nil {
			return nil, nil, nil, fmt.Errorf("error writing stl: %v", err)
		}
		generatedFiles = append(generatedFiles, outputFilename)
	}

	if wantsType("svg") {
		outputFilename := baseName + ".svg"
		if err := WriteSVG(outputFilename, gf, &bounds); err != nil {
			return nil, nil, nil, fmt.Errorf("error writing svg: %v", err)
		}
		generatedFiles = append(generatedFiles, outputFilename)
	}

	if wantsType("png") {
		outputFilename := baseName + ".png"
		if f, err := os.Create(outputFilename); err == nil {
			png.Encode(f, img)
			f.Close()
			generatedFiles = append(generatedFiles, outputFilename)
		}
	}

	if wantsType("scad") {
		outputFilename := baseName + ".scad"
		if err := WriteStencilSCAD(outputFilename, gf, outlineGf, cfg, &bounds); err != nil {
			return nil, nil, nil, fmt.Errorf("error writing scad: %v", err)
		}
		generatedFiles = append(generatedFiles, outputFilename)
	}

	return generatedFiles, img, outlineImg, nil
}