package main

import (
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"os"
	"path/filepath"
)

// ---------------------------------------------------------------------------
// Palette
// ---------------------------------------------------------------------------

type rgb struct{ R, G, B uint8 }

func (c rgb) nrgba() color.RGBA { return color.RGBA{c.R, c.G, c.B, 0xff} }

// ansi16 is the standard xterm palette for SGR 30-37 / 90-97.
var ansi16 = [16]rgb{
	{0, 0, 0}, {205, 0, 0}, {0, 205, 0}, {205, 205, 0},
	{0, 0, 238}, {205, 0, 205}, {0, 205, 205}, {229, 229, 229},
	{127, 127, 127}, {255, 0, 0}, {0, 255, 0}, {255, 255, 0},
	{92, 92, 255}, {255, 0, 255}, {0, 255, 255}, {255, 255, 255},
}

// cubeLevels are the six intensity steps of the 216-colour cube (indices 16-231).
var cubeLevels = [6]uint8{0, 95, 135, 175, 215, 255}

// palette256 resolves an xterm 256-colour index to RGB.
func palette256(i uint8) rgb {
	switch {
	case i < 16:
		return ansi16[i]
	case i < 232:
		n := int(i) - 16
		return rgb{cubeLevels[n/36], cubeLevels[(n/6)%6], cubeLevels[n%6]}
	default:
		v := uint8(8 + 10*(int(i)-232))
		return rgb{v, v, v}
	}
}

// theme carries the default foreground/background and the chrome colours used
// by the contact sheet.
type theme struct {
	Name    string
	FG      rgb
	BG      rgb
	Chrome  rgb // sheet background
	Label   rgb // sheet label text
	Border  rgb // thumbnail border
	IsLight bool
}

var themeDark = theme{
	Name: "dark",
	FG:   rgb{222, 222, 226}, BG: rgb{18, 18, 22},
	Chrome: rgb{34, 34, 40}, Label: rgb{215, 215, 220}, Border: rgb{80, 80, 90},
}

var themeLight = theme{
	Name: "light",
	FG:   rgb{32, 32, 38}, BG: rgb{250, 250, 247},
	Chrome: rgb{226, 226, 222}, Label: rgb{40, 40, 46}, Border: rgb{150, 150, 148},
	IsLight: true,
}

func lookupTheme(name string) (theme, error) {
	switch name {
	case "dark", "":
		return themeDark, nil
	case "light":
		return themeLight, nil
	}
	return theme{}, fmt.Errorf("unknown theme %q (want dark or light)", name)
}

// resolve turns a termColor into concrete RGB against a theme.
func (th theme) resolve(c termColor, isFG bool) rgb {
	switch c.Mode {
	case colIndexed:
		return palette256(c.Idx)
	case colRGB:
		return rgb{c.R, c.G, c.B}
	}
	if isFG {
		return th.FG
	}
	return th.BG
}

func mix(a, b rgb, f float64) rgb {
	if f < 0 {
		f = 0
	}
	if f > 1 {
		f = 1
	}
	m := func(x, y uint8) uint8 { return uint8(float64(x)*(1-f) + float64(y)*f + 0.5) }
	return rgb{m(a.R, b.R), m(a.G, b.G), m(a.B, b.B)}
}

// cellColors applies bold, dim and reverse to produce the pair actually painted.
//
//	bold     brightens: an ANSI colour 0-7 becomes its 8-15 bright twin, any
//	         other colour is pushed 25% toward white
//	dim      blends the foreground halfway into the background
//	reverse  swaps foreground and background AFTER the two above
func (th theme) cellColors(c cell) (fg, bg rgb) {
	fg = th.resolve(c.FG, true)
	bg = th.resolve(c.BG, false)
	if c.Attrs&attrBold != 0 {
		if c.FG.Mode == colIndexed && c.FG.Idx < 8 {
			fg = ansi16[c.FG.Idx+8]
		} else {
			fg = mix(fg, rgb{255, 255, 255}, 0.25)
		}
	}
	if c.Attrs&attrDim != 0 {
		fg = mix(fg, bg, 0.5)
	}
	if c.Attrs&attrReverse != 0 {
		fg, bg = bg, fg
	}
	return fg, bg
}

// ---------------------------------------------------------------------------
// Size guards
// ---------------------------------------------------------------------------

const (
	maxScale  = 32
	maxPad    = 400
	maxDim    = 20000
	maxPixels = 64 << 20 // 64M pixels; an RGBA buffer of that is 256 MiB
)

func checkImageSize(w, h int, what string) error {
	if w <= 0 || h <= 0 {
		return fmt.Errorf("%s would be %dx%d pixels, which is empty", what, w, h)
	}
	if w > maxDim || h > maxDim {
		return fmt.Errorf("%s would be %dx%d pixels; the per-side limit is %d - reduce --scale, --cols or --rows",
			what, w, h, maxDim)
	}
	if int64(w)*int64(h) > maxPixels {
		return fmt.Errorf("%s would be %dx%d = %d pixels (%s of RGBA); the limit is %d pixels - reduce --scale, --cols or --rows",
			what, w, h, int64(w)*int64(h), humanBytes(int64(w)*int64(h)*4), maxPixels)
	}
	return nil
}

// ---------------------------------------------------------------------------
// Grid rendering
// ---------------------------------------------------------------------------

// renderOpts controls how a grid becomes pixels.
type renderOpts struct {
	Scale int   // integer nearest-neighbour magnification, >= 1
	Pad   int   // border in unscaled pixels
	Theme theme // palette for default colours
}

func (o renderOpts) validate() error {
	if o.Scale < 1 || o.Scale > maxScale {
		return fmt.Errorf("--scale must be between 1 and %d, got %d", maxScale, o.Scale)
	}
	if o.Pad < 0 || o.Pad > maxPad {
		return fmt.Errorf("--pad must be between 0 and %d, got %d", maxPad, o.Pad)
	}
	return nil
}

// gridPixelSize is the unscaled size of a grid image including padding.
func gridPixelSize(cols, rows, pad int) (int, int) {
	return cols*cellW + 2*pad, rows*cellH + 2*pad
}

// renderGrid paints the terminal grid into a fresh RGBA image.
func renderGrid(t *terminal, o renderOpts) (*image.RGBA, error) {
	if err := o.validate(); err != nil {
		return nil, err
	}
	w0, h0 := gridPixelSize(t.cols, t.rows, o.Pad)
	w, h := w0*o.Scale, h0*o.Scale
	if err := checkImageSize(w, h, "the frame"); err != nil {
		return nil, err
	}
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	draw.Draw(img, img.Bounds(), &image.Uniform{o.Theme.BG.nrgba()}, image.Point{}, draw.Src)

	s := o.Scale
	for y := 0; y < t.rows; y++ {
		for x := 0; x < t.cols; x++ {
			c := *t.at(y, x)
			fg, bg := o.Theme.cellColors(c)
			ox := (o.Pad + x*cellW) * s
			oy := (o.Pad + y*cellH) * s
			if bg != o.Theme.BG {
				draw.Draw(img, image.Rect(ox, oy, ox+cellW*s, oy+cellH*s),
					&image.Uniform{bg.nrgba()}, image.Point{}, draw.Src)
			}
			drawGlyph(img, ox+glyphX*s, oy+glyphY*s, s, c.Ch, fg)
			if c.Attrs&attrUnderline != 0 {
				uy := oy + (cellH-2)*s
				draw.Draw(img, image.Rect(ox, uy, ox+cellW*s, uy+s),
					&image.Uniform{fg.nrgba()}, image.Point{}, draw.Src)
			}
		}
	}
	return img, nil
}

// glyphRows returns the packed scan lines for a rune, falling back to the
// missing-glyph box for anything outside printable ASCII.
func glyphRows(r rune) [8]uint8 {
	if r >= 0x20 && r <= 0x7e {
		return fontBits[r-0x20]
	}
	if r == ' ' {
		return fontBits[0]
	}
	return fontFallback
}

// drawGlyph blits one character. Each font pixel becomes a scale x scale block.
func drawGlyph(img *image.RGBA, ox, oy, scale int, r rune, fg rgb) {
	if r == ' ' || r == 0 {
		return
	}
	rows := glyphRows(r)
	col := fg.nrgba()
	for gy := 0; gy < glyphH; gy++ {
		bits := rows[gy]
		if bits == 0 {
			continue
		}
		for gx := 0; gx < glyphW; gx++ {
			if bits&(1<<uint(glyphW-1-gx)) == 0 {
				continue
			}
			px := ox + gx*scale
			py := oy + gy*scale
			for dy := 0; dy < scale; dy++ {
				for dx := 0; dx < scale; dx++ {
					img.SetRGBA(px+dx, py+dy, col)
				}
			}
		}
	}
}

// drawText writes a line of ASCII at scale 1 (used for contact-sheet chrome).
func drawText(img *image.RGBA, x, y int, s string, fg rgb) {
	for i, r := range s {
		drawGlyph(img, x+i*cellW, y, 1, r, fg)
	}
}

// ---------------------------------------------------------------------------
// Downscaling (contact-sheet thumbnails)
// ---------------------------------------------------------------------------

// downscale box-averages src into a dstW x dstH image. Averaging rather than
// point-sampling matters here: a 1-pixel-wide glyph stroke would vanish
// entirely under nearest-neighbour, leaving blank-looking thumbnails.
func downscale(src *image.RGBA, dstW, dstH int) *image.RGBA {
	sb := src.Bounds()
	sw, sh := sb.Dx(), sb.Dy()
	if dstW < 1 {
		dstW = 1
	}
	if dstH < 1 {
		dstH = 1
	}
	dst := image.NewRGBA(image.Rect(0, 0, dstW, dstH))
	for dy := 0; dy < dstH; dy++ {
		y0 := dy * sh / dstH
		y1 := (dy + 1) * sh / dstH
		if y1 <= y0 {
			y1 = y0 + 1
		}
		for dx := 0; dx < dstW; dx++ {
			x0 := dx * sw / dstW
			x1 := (dx + 1) * sw / dstW
			if x1 <= x0 {
				x1 = x0 + 1
			}
			var sr, sg, sb2, n uint32
			for y := y0; y < y1; y++ {
				for x := x0; x < x1; x++ {
					c := src.RGBAAt(sb.Min.X+x, sb.Min.Y+y)
					sr += uint32(c.R)
					sg += uint32(c.G)
					sb2 += uint32(c.B)
					n++
				}
			}
			if n == 0 {
				n = 1
			}
			dst.SetRGBA(dx, dy, color.RGBA{uint8(sr / n), uint8(sg / n), uint8(sb2 / n), 0xff})
		}
	}
	return dst
}

// ---------------------------------------------------------------------------
// Output files
// ---------------------------------------------------------------------------

// writePNG encodes img to path. An existing file is never clobbered unless
// force is set, and the encode goes to a temporary file first so a failure
// halfway through cannot leave a truncated PNG behind.
func writePNG(path string, img image.Image, force bool) error {
	if _, err := os.Stat(path); err == nil {
		if !force {
			return fmt.Errorf("%s already exists (pass --force to overwrite)", path)
		}
	} else if !os.IsNotExist(err) {
		return fmt.Errorf("cannot check %s: %w", path, err)
	}
	dir := filepath.Dir(path)
	if dir != "" && dir != "." {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return fmt.Errorf("cannot create %s: %w", dir, err)
		}
	}
	tmp, err := os.CreateTemp(dir, ".capturestudio-*.png")
	if err != nil {
		return fmt.Errorf("cannot create a temporary file in %s: %w", dir, err)
	}
	tmpName := tmp.Name()
	if err := png.Encode(tmp, img); err != nil {
		tmp.Close()
		os.Remove(tmpName)
		return fmt.Errorf("cannot encode %s: %w", path, err)
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmpName)
		return fmt.Errorf("cannot write %s: %w", path, err)
	}
	if err := os.Rename(tmpName, path); err != nil {
		os.Remove(tmpName)
		return fmt.Errorf("cannot write %s: %w", path, err)
	}
	return nil
}

func fileSize(path string) int64 {
	if fi, err := os.Stat(path); err == nil {
		return fi.Size()
	}
	return 0
}