package main

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
)

// Standard cylindrical cell dimensions: diameter × length (mm)
var StandardCells = map[string][2]float64{
	"10440": {10, 44},
	"14500": {14, 50},
	"16340": {16, 34},
	"18350": {18, 35},
	"18650": {18, 65},
	"20700": {20, 70},
	"21700": {21, 70},
	"26650": {26, 65},
	"32650": {32, 65},
}

// BatteryConfig defines a battery compartment to be generated inside an enclosure.
type BatteryConfig struct {
	CellSize     string               `json:"cellSize"`     // preset key or "custom"
	CellDiameter float64              `json:"cellDiameter"` // mm, used when CellSize == "custom"
	CellLength   float64              `json:"cellLength"`   // mm, used when CellSize == "custom"
	Series       int                  `json:"series"`
	Parallel     int                  `json:"parallel"`
	PositionX    float64              `json:"positionX"`   // bay origin X in enclosure coords
	PositionY    float64              `json:"positionY"`   // bay origin Y in enclosure coords
	Rotation     float64              `json:"rotation"`    // degrees, placement rotation on the canvas
	Placement    string               `json:"placement"`   // "bottom", "side-1", "side-2", etc.
	Orientation  string               `json:"orientation"` // "horizontal" or "vertical"
	WallThick    float64              `json:"wallThick"`   // internal wall thickness (default 1.2)
	Clearance    float64              `json:"clearance"`   // gap around each cell (default 0.3)
	HasLid       bool                 `json:"hasLid"`      // generate a clip-on lid for the bay
	RoutingHoles []BatteryRoutingHole `json:"routingHoles"`
}

// BatteryRoutingHole is a wire passthrough in the bay's outer walls or floor
// to route power wires from the battery bay to the PCB.
type BatteryRoutingHole struct {
	Face     string  `json:"face"`     // "front", "back", "left", "right", "floor"
	Position float64 `json:"position"` // 0..1 along the face width
	Height   float64 `json:"height"`   // 0..1 along the face height (ignored for floor)
	Diameter float64 `json:"diameter"` // hole diameter mm
}

// BatteryLayout holds computed dimensions for a battery bay.
type BatteryLayout struct {
	BayWidth   float64 `json:"bayWidth"`
	BayDepth   float64 `json:"bayDepth"`
	BayHeight  float64 `json:"bayHeight"`
	Rows       int     `json:"rows"`
	Cols       int     `json:"cols"`
	CellDia    float64 `json:"cellDia"`
	CellLen    float64 `json:"cellLen"`
	Wall       float64 `json:"wall"`
	Clr        float64 `json:"clr"`
	TotalCells int     `json:"totalCells"`
}

func (cfg *BatteryConfig) resolveDimensions() (dia, length float64) {
	if dims, ok := StandardCells[cfg.CellSize]; ok {
		return dims[0], dims[1]
	}
	return cfg.CellDiameter, cfg.CellLength
}

func (cfg *BatteryConfig) wallThick() float64 {
	if cfg.WallThick > 0 {
		return cfg.WallThick
	}
	return 1.2
}

func (cfg *BatteryConfig) clearance() float64 {
	if cfg.Clearance > 0 {
		return cfg.Clearance
	}
	return 0.3
}

// ComputeLayout calculates the physical dimensions of the battery bay.
func (cfg *BatteryConfig) ComputeLayout() BatteryLayout {
	dia, length := cfg.resolveDimensions()
	wall := cfg.wallThick()
	clr := cfg.clearance()

	rows := cfg.Series
	cols := cfg.Parallel
	if rows < 1 {
		rows = 1
	}
	if cols < 1 {
		cols = 1
	}

	cellSlot := dia + 2*clr

	var bayW, bayD, bayH float64

	if cfg.Orientation == "vertical" {
		bayW = float64(cols)*(cellSlot+wall) + wall
		bayD = float64(rows)*(cellSlot+wall) + wall
		bayH = length + clr + wall
	} else {
		// horizontal: cell axis along X
		bayW = length + 2*clr + 2*wall
		bayD = float64(cols)*(cellSlot+wall) + wall
		bayH = float64(rows)*(cellSlot+wall) + wall
	}

	return BatteryLayout{
		BayWidth:   bayW,
		BayDepth:   bayD,
		BayHeight:  bayH,
		Rows:       rows,
		Cols:       cols,
		CellDia:    dia,
		CellLen:    length,
		Wall:       wall,
		Clr:        clr,
		TotalCells: rows * cols,
	}
}

// WriteBatteryBaySCAD writes the battery bay base piece to a file.
func WriteBatteryBaySCAD(filename string, cfg *BatteryConfig) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	return writeBatteryBaySCADTo(f, cfg)
}

// GenerateBatteryBaySCADString returns the bay SCAD as a string.
func GenerateBatteryBaySCADString(cfg *BatteryConfig) (string, error) {
	var buf bytes.Buffer
	if err := writeBatteryBaySCADTo(&buf, cfg); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func writeBatteryBaySCADTo(f io.Writer, cfg *BatteryConfig) error {
	lay := cfg.ComputeLayout()
	wall := lay.Wall
	clr := lay.Clr
	cellSlot := lay.CellDia + 2*clr

	fmt.Fprintf(f, "// Battery bay — %ds%dp %s cells (%s)\n", lay.Rows, lay.Cols, cfg.CellSize, cfg.Orientation)
	fmt.Fprintf(f, "$fn = 48;\n\n")
	fmt.Fprintf(f, "cell_dia = %f;\n", lay.CellDia)
	fmt.Fprintf(f, "cell_len = %f;\n", lay.CellLen)
	fmt.Fprintf(f, "wall = %f;\n", wall)
	fmt.Fprintf(f, "clr = %f;\n", clr)
	fmt.Fprintf(f, "cell_slot = cell_dia + 2*clr;\n")
	fmt.Fprintf(f, "bay_w = %f;\n", lay.BayWidth)
	fmt.Fprintf(f, "bay_d = %f;\n", lay.BayDepth)
	fmt.Fprintf(f, "bay_h = %f;\n\n", lay.BayHeight)

	if cfg.Orientation == "vertical" {
		writeVerticalBay(f, lay, cellSlot, wall, clr, cfg)
	} else {
		writeHorizontalBay(f, lay, cellSlot, wall, clr, cfg)
	}

	return nil
}

func writeHorizontalBay(f io.Writer, lay BatteryLayout, cellSlot, wall, clr float64, cfg *BatteryConfig) {
	cellCavityLen := lay.CellLen + 2*clr

	fmt.Fprintf(f, "// Battery bay base — cells lie horizontal, open top for insertion\n")
	fmt.Fprintf(f, "difference() {\n")
	fmt.Fprintf(f, "  cube([bay_w, bay_d, bay_h]);\n\n")

	for row := 0; row < lay.Rows; row++ {
		for col := 0; col < lay.Cols; col++ {
			cy := wall + cellSlot/2 + float64(col)*(cellSlot+wall)
			cz := wall + cellSlot/2 + float64(row)*(cellSlot+wall)

			fmt.Fprintf(f, "  // Cell [%d,%d]\n", row, col)
			fmt.Fprintf(f, "  translate([wall, %f, %f])\n", cy, cz)
			fmt.Fprintf(f, "    rotate([0, 90, 0])\n")
			fmt.Fprintf(f, "    cylinder(d=cell_slot, h=%f);\n", cellCavityLen)

			// Open above cell center for top insertion
			fmt.Fprintf(f, "  translate([wall, %f, %f])\n", cy-cellSlot/2, cz)
			fmt.Fprintf(f, "    cube([%f, cell_slot, bay_h]);\n\n", cellCavityLen)
		}
	}

	writeRoutingHoles(f, lay, cfg)

	fmt.Fprintf(f, "}\n")
}

func writeVerticalBay(f io.Writer, lay BatteryLayout, cellSlot, wall, clr float64, cfg *BatteryConfig) {
	fmt.Fprintf(f, "// Battery bay base — cells stand vertical\n")
	fmt.Fprintf(f, "difference() {\n")
	fmt.Fprintf(f, "  cube([bay_w, bay_d, bay_h]);\n\n")

	for row := 0; row < lay.Rows; row++ {
		for col := 0; col < lay.Cols; col++ {
			cx := wall + cellSlot/2 + float64(col)*(cellSlot+wall)
			cy := wall + cellSlot/2 + float64(row)*(cellSlot+wall)

			fmt.Fprintf(f, "  // Cell [%d,%d]\n", row, col)
			fmt.Fprintf(f, "  translate([%f, %f, wall])\n", cx, cy)
			fmt.Fprintf(f, "    cylinder(d=cell_slot, h=bay_h);\n\n")
		}
	}

	writeRoutingHoles(f, lay, cfg)

	fmt.Fprintf(f, "}\n")
}

func writeRoutingHoles(f io.Writer, lay BatteryLayout, cfg *BatteryConfig) {
	wall := lay.Wall
	for _, rh := range cfg.RoutingHoles {
		d := rh.Diameter
		if d <= 0 {
			d = 3.0
		}
		pos := rh.Position
		if pos < 0 {
			pos = 0
		}
		if pos > 1 {
			pos = 1
		}
		ht := rh.Height
		if ht < 0 {
			ht = 0
		}
		if ht > 1 {
			ht = 1
		}

		switch rh.Face {
		case "front":
			cx := wall + pos*(lay.BayWidth-2*wall)
			cz := wall + ht*(lay.BayHeight-2*wall)
			fmt.Fprintf(f, "  // Routing hole (front wall)\n")
			fmt.Fprintf(f, "  translate([%f, -0.1, %f])\n", cx, cz)
			fmt.Fprintf(f, "    rotate([-90, 0, 0])\n")
			fmt.Fprintf(f, "    cylinder(d=%f, h=%f);\n\n", d, wall+0.2)
		case "back":
			cx := wall + pos*(lay.BayWidth-2*wall)
			cz := wall + ht*(lay.BayHeight-2*wall)
			fmt.Fprintf(f, "  // Routing hole (back wall)\n")
			fmt.Fprintf(f, "  translate([%f, %f, %f])\n", cx, lay.BayDepth-wall-0.1, cz)
			fmt.Fprintf(f, "    rotate([-90, 0, 0])\n")
			fmt.Fprintf(f, "    cylinder(d=%f, h=%f);\n\n", d, wall+0.2)
		case "left":
			cy := wall + pos*(lay.BayDepth-2*wall)
			cz := wall + ht*(lay.BayHeight-2*wall)
			fmt.Fprintf(f, "  // Routing hole (left wall)\n")
			fmt.Fprintf(f, "  translate([-0.1, %f, %f])\n", cy, cz)
			fmt.Fprintf(f, "    rotate([0, 90, 0])\n")
			fmt.Fprintf(f, "    cylinder(d=%f, h=%f);\n\n", d, wall+0.2)
		case "right":
			cy := wall + pos*(lay.BayDepth-2*wall)
			cz := wall + ht*(lay.BayHeight-2*wall)
			fmt.Fprintf(f, "  // Routing hole (right wall)\n")
			fmt.Fprintf(f, "  translate([%f, %f, %f])\n", lay.BayWidth-wall-0.1, cy, cz)
			fmt.Fprintf(f, "    rotate([0, 90, 0])\n")
			fmt.Fprintf(f, "    cylinder(d=%f, h=%f);\n\n", d, wall+0.2)
		case "floor":
			cx := wall + pos*(lay.BayWidth-2*wall)
			cy := wall + ht*(lay.BayDepth-2*wall)
			fmt.Fprintf(f, "  // Routing hole (floor)\n")
			fmt.Fprintf(f, "  translate([%f, %f, -0.1])\n", cx, cy)
			fmt.Fprintf(f, "    cylinder(d=%f, h=%f);\n\n", d, wall+0.2)
		}
	}
}

// WriteBatterySeparatorSCAD writes the egg-carton separator piece.
func WriteBatterySeparatorSCAD(filename string, cfg *BatteryConfig) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	return writeBatterySeparatorSCADTo(f, cfg)
}

// GenerateBatterySeparatorSCADString returns the separator SCAD as a string.
func GenerateBatterySeparatorSCADString(cfg *BatteryConfig) (string, error) {
	var buf bytes.Buffer
	if err := writeBatterySeparatorSCADTo(&buf, cfg); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func writeBatterySeparatorSCADTo(f io.Writer, cfg *BatteryConfig) error {
	lay := cfg.ComputeLayout()
	wall := lay.Wall
	clr := lay.Clr
	cellSlot := lay.CellDia + 2*clr

	if cfg.Orientation == "vertical" {
		return writeVerticalSeparator(f, lay, cellSlot, wall, clr)
	}
	return writeHorizontalSeparator(f, lay, cellSlot, wall, clr)
}

func writeHorizontalSeparator(f io.Writer, lay BatteryLayout, cellSlot, wall, clr float64) error {
	// The separator mirrors the bay base: same external dimensions, with
	// egg-carton bumps that capture the top half of each cell.
	// Total height = same as bay so it registers properly.
	cellCavityLen := lay.CellLen + 2*clr

	fmt.Fprintf(f, "// Battery separator (egg-carton retainer) — horizontal cells\n")
	fmt.Fprintf(f, "$fn = 48;\n\n")
	fmt.Fprintf(f, "cell_slot = %f;\n", cellSlot)
	fmt.Fprintf(f, "wall = %f;\n", wall)
	fmt.Fprintf(f, "bay_w = %f;\n", lay.BayWidth)
	fmt.Fprintf(f, "bay_d = %f;\n", lay.BayDepth)
	fmt.Fprintf(f, "bay_h = %f;\n\n", lay.BayHeight)

	fmt.Fprintf(f, "difference() {\n")
	fmt.Fprintf(f, "  cube([bay_w, bay_d, bay_h]);\n\n")

	for row := 0; row < lay.Rows; row++ {
		for col := 0; col < lay.Cols; col++ {
			cy := wall + cellSlot/2 + float64(col)*(cellSlot+wall)
			cz := wall + cellSlot/2 + float64(row)*(cellSlot+wall)

			fmt.Fprintf(f, "  // Cell [%d,%d]\n", row, col)
			// Cell cylinder — same position as bay
			fmt.Fprintf(f, "  translate([wall, %f, %f])\n", cy, cz)
			fmt.Fprintf(f, "    rotate([0, 90, 0])\n")
			fmt.Fprintf(f, "    cylinder(d=cell_slot, h=%f);\n", cellCavityLen)

			// Cut below cell center (keep top half = retainer bump)
			fmt.Fprintf(f, "  translate([wall, %f, -0.1])\n", cy-cellSlot/2)
			fmt.Fprintf(f, "    cube([%f, cell_slot, %f]);\n\n", cellCavityLen, cz+0.1)
		}
	}

	fmt.Fprintf(f, "}\n")
	return nil
}

func writeVerticalSeparator(f io.Writer, lay BatteryLayout, cellSlot, wall, clr float64) error {
	// For vertical cells, the separator is a plate with circular holes
	// that sit over the cells, providing lateral retention.
	plateH := wall
	lipDepth := cellSlot * 0.3 // how far the lip extends down into the cell socket
	sepH := plateH + lipDepth

	fmt.Fprintf(f, "// Battery separator (retainer plate) — vertical cells\n")
	fmt.Fprintf(f, "$fn = 48;\n\n")

	fmt.Fprintf(f, "difference() {\n")
	fmt.Fprintf(f, "  cube([%f, %f, %f]);\n\n", lay.BayWidth, lay.BayDepth, sepH)

	for row := 0; row < lay.Rows; row++ {
		for col := 0; col < lay.Cols; col++ {
			cx := wall + cellSlot/2 + float64(col)*(cellSlot+wall)
			cy := wall + cellSlot/2 + float64(row)*(cellSlot+wall)

			// Hole slightly smaller than cell so it retains but allows insertion
			retainDia := lay.CellDia - 1.0
			if retainDia < lay.CellDia*0.7 {
				retainDia = lay.CellDia * 0.7
			}
			fmt.Fprintf(f, "  // Cell [%d,%d] retention hole\n", row, col)
			fmt.Fprintf(f, "  translate([%f, %f, -0.1])\n", cx, cy)
			fmt.Fprintf(f, "    cylinder(d=%f, h=%f);\n\n", retainDia, sepH+0.2)
		}
	}

	fmt.Fprintf(f, "}\n")
	return nil
}

// WriteBatteryLidSCAD writes the clip-on lid for the battery bay.
func WriteBatteryLidSCAD(filename string, cfg *BatteryConfig) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	return writeBatteryLidSCADTo(f, cfg)
}

// GenerateBatteryLidSCADString returns the lid SCAD as a string.
func GenerateBatteryLidSCADString(cfg *BatteryConfig) (string, error) {
	var buf bytes.Buffer
	if err := writeBatteryLidSCADTo(&buf, cfg); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func writeBatteryLidSCADTo(f io.Writer, cfg *BatteryConfig) error {
	lay := cfg.ComputeLayout()
	wall := lay.Wall
	lipH := 2.0      // how far the lid lip extends down into the bay
	lidPlate := wall // top plate thickness
	clipTol := 0.15  // tolerance for clip fit

	fmt.Fprintf(f, "// Battery bay lid — clips onto bay walls\n")
	fmt.Fprintf(f, "$fn = 48;\n\n")

	outerW := lay.BayWidth + 2*clipTol
	outerD := lay.BayDepth + 2*clipTol
	innerW := lay.BayWidth - 2*wall + 2*clipTol
	innerD := lay.BayDepth - 2*wall + 2*clipTol

	fmt.Fprintf(f, "difference() {\n")
	fmt.Fprintf(f, "  cube([%f, %f, %f]);\n", outerW, outerD, lipH+lidPlate)
	fmt.Fprintf(f, "  translate([%f, %f, -0.1])\n", (outerW-innerW)/2, (outerD-innerD)/2)
	fmt.Fprintf(f, "    cube([%f, %f, %f]);\n", innerW, innerD, lipH+0.1)
	fmt.Fprintf(f, "}\n")

	return nil
}

// GenerateBatteryOutputs generates all battery SCAD files for the given config.
func GenerateBatteryOutputs(cfg *BatteryConfig, outputDir string) ([]string, error) {
	if cfg == nil {
		return nil, fmt.Errorf("no battery config")
	}
	os.MkdirAll(outputDir, 0755)

	var files []string

	bayPath := filepath.Join(outputDir, "battery_bay.scad")
	if err := WriteBatteryBaySCAD(bayPath, cfg); err != nil {
		return nil, fmt.Errorf("write battery bay: %v", err)
	}
	files = append(files, bayPath)

	if cfg.HasLid {
		lidPath := filepath.Join(outputDir, "battery_lid.scad")
		if err := WriteBatteryLidSCAD(lidPath, cfg); err != nil {
			return nil, fmt.Errorf("write battery lid: %v", err)
		}
		files = append(files, lidPath)
	} else {
		sepPath := filepath.Join(outputDir, "battery_separator.scad")
		if err := WriteBatterySeparatorSCAD(sepPath, cfg); err != nil {
			return nil, fmt.Errorf("write battery separator: %v", err)
		}
		files = append(files, sepPath)
	}

	return files, nil
}