package main

import (
	"crypto/rand"
	"embed"
	"encoding/binary"
	"encoding/hex"
	"flag"
	"fmt"
	"html/template"
	"image"
	"image/png"
	"io"
	"log"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// --- Configuration ---

type Config struct {
	StencilHeight float64
	WallHeight    float64
	WallThickness float64
	DPI           float64
	KeepPNG       bool
}

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

// --- 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, pixelToMM 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, 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 wallMask []bool
	var boardMask []bool
	if outlineImg != nil {
		fmt.Println("Computing wall mask...")
		wallMask, boardMask = ComputeWallMask(outlineImg, cfg.WallThickness, pixelToMM)
	}

	// 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 = cfg.WallHeight
			} else if isStencilSolid {
				if isInsideBoard {
					h = cfg.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
}

// --- Logic ---

func processPCB(gerberPath, outlinePath string, cfg Config) (string, error) {
	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 {
		return "", 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 "", fmt.Errorf("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
	margin := cfg.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(cfg.DPI, &bounds)

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

	if cfg.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()
		}
	}

	// 4. Generate Mesh
	fmt.Println("Generating mesh...")
	triangles := GenerateMeshFromImages(img, outlineImg, cfg)

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

	return outputPath, nil
}

// --- CLI ---

func runCLI(cfg Config, args []string) {
	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]
	}

	_, err := processPCB(gerberPath, outlinePath, cfg)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}
	fmt.Println("Success! Happy printing.")
}

// --- Server ---

//go:embed static/*
var staticFiles embed.FS

func randomID() string {
	b := make([]byte, 16)
	rand.Read(b)
	return hex.EncodeToString(b)
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}
	// Read index.html from embedded FS
	content, err := staticFiles.ReadFile("static/index.html")
	if err != nil {
		http.Error(w, "Could not load index page", http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "text/html")
	w.Write(content)
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Create temp dir
	tempDir := filepath.Join(".", "temp")
	os.MkdirAll(tempDir, 0755)

	uuid := randomID()

	// Parse params
	height, _ := strconv.ParseFloat(r.FormValue("height"), 64)
	dpi, _ := strconv.ParseFloat(r.FormValue("dpi"), 64)
	wallHeight, _ := strconv.ParseFloat(r.FormValue("wallHeight"), 64)
	wallThickness, _ := strconv.ParseFloat(r.FormValue("wallThickness"), 64)

	if height == 0 {
		height = DefaultStencilHeight
	}
	if dpi == 0 {
		dpi = DefaultDPI
	}
	if wallHeight == 0 {
		wallHeight = DefaultWallHeight
	}
	if wallThickness == 0 {
		wallThickness = DefaultWallThickness
	}

	cfg := Config{
		StencilHeight: height,
		WallHeight:    wallHeight,
		WallThickness: wallThickness,
		DPI:           dpi,
		KeepPNG:       false,
	}

	// Handle Gerber File
	file, header, err := r.FormFile("gerber")
	if err != nil {
		http.Error(w, "Error retrieving gerber file", http.StatusBadRequest)
		return
	}
	defer file.Close()

	gerberPath := filepath.Join(tempDir, uuid+"_paste"+filepath.Ext(header.Filename))
	outFile, err := os.Create(gerberPath)
	if err != nil {
		http.Error(w, "Server error saving file", http.StatusInternalServerError)
		return
	}
	defer outFile.Close()
	io.Copy(outFile, file)

	// Handle Outline File (Optional)
	outlineFile, outlineHeader, err := r.FormFile("outline")
	var outlinePath string
	if err == nil {
		defer outlineFile.Close()
		outlinePath = filepath.Join(tempDir, uuid+"_outline"+filepath.Ext(outlineHeader.Filename))
		outOutline, err := os.Create(outlinePath)
		if err == nil {
			defer outOutline.Close()
			io.Copy(outOutline, outlineFile)
		}
	}

	// Process
	outSTL, err := processPCB(gerberPath, outlinePath, cfg)
	if err != nil {
		log.Printf("Error processing: %v", err)
		http.Error(w, fmt.Sprintf("Error processing PCB: %v", err), http.StatusInternalServerError)
		return
	}

	// Render Success
	tmpl, err := template.ParseFS(staticFiles, "static/result.html")
	if err != nil {
		http.Error(w, "Template error", http.StatusInternalServerError)
		return
	}
	data := struct{ Filename string }{Filename: filepath.Base(outSTL)}
	tmpl.Execute(w, data)
}

func downloadHandler(w http.ResponseWriter, r *http.Request) {
	vars := strings.Split(r.URL.Path, "/")
	if len(vars) < 3 {
		http.NotFound(w, r)
		return
	}
	filename := vars[2]

	// Security check: ensure no path traversal
	if strings.Contains(filename, "/") || strings.Contains(filename, "\\") || strings.Contains(filename, "..") {
		http.Error(w, "Invalid filename", http.StatusBadRequest)
		return
	}

	path := filepath.Join("temp", filename)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		http.NotFound(w, r)
		return
	}

	w.Header().Set("Content-Disposition", "attachment; filename="+filename)
	w.Header().Set("Content-Type", "application/octet-stream")
	http.ServeFile(w, r, path)
}

func runServer(port string) {
	// Serve static files (CSS, etc.)
	// This will serve files under /static/ from the embedded fs
	http.Handle("/static/", http.FileServer(http.FS(staticFiles)))

	http.HandleFunc("/", indexHandler)
	http.HandleFunc("/upload", uploadHandler)
	http.HandleFunc("/download/", downloadHandler)

	fmt.Printf("Starting server on http://0.0.0.0:%s\n", port)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}

// --- Main ---

var (
	flagStencilHeight float64
	flagWallHeight    float64
	flagWallThickness float64
	flagDPI           float64
	flagKeepPNG       bool
	flagServer        bool
	flagPort          string
)

func main() {
	flag.Float64Var(&flagStencilHeight, "height", DefaultStencilHeight, "Stencil height in mm")
	flag.Float64Var(&flagWallHeight, "wall-height", DefaultWallHeight, "Wall height in mm")
	flag.Float64Var(&flagWallThickness, "wall-thickness", DefaultWallThickness, "Wall thickness in mm")
	flag.Float64Var(&flagDPI, "dpi", DefaultDPI, "DPI for rendering (lower = smaller file, rougher curves)")
	flag.BoolVar(&flagKeepPNG, "keep-png", false, "Save intermediate PNG file")

	flag.BoolVar(&flagServer, "server", false, "Start in server mode")
	flag.StringVar(&flagPort, "port", "8080", "Port to run the server on")

	flag.Parse()

	if flagServer {
		runServer(flagPort)
	} else {
		cfg := Config{
			StencilHeight: flagStencilHeight,
			WallHeight:    flagWallHeight,
			WallThickness: flagWallThickness,
			DPI:           flagDPI,
			KeepPNG:       flagKeepPNG,
		}
		runCLI(cfg, flag.Args())
	}
}