package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"math"
	"os"
	"strconv"
	"strings"
)

// ---------------------------------------------------------------------------
// Colours
// ---------------------------------------------------------------------------

const (
	defaultCircleRadius = 20
	strokeWidth         = 3.0 // px, for rectangles and arrow shafts
)

var namedColors = map[string]color.NRGBA{
	"red":     {0xD9, 0x30, 0x25, 0xFF},
	"orange":  {0xE8, 0x71, 0x0A, 0xFF},
	"yellow":  {0xF2, 0xB6, 0x00, 0xFF},
	"green":   {0x18, 0x80, 0x38, 0xFF},
	"blue":    {0x17, 0x57, 0xC0, 0xFF},
	"purple":  {0x7B, 0x1F, 0xA2, 0xFF},
	"magenta": {0xC2, 0x18, 0x5B, 0xFF},
	"black":   {0x11, 0x11, 0x11, 0xFF},
	"white":   {0xFF, 0xFF, 0xFF, 0xFF},
	"grey":    {0x5F, 0x63, 0x68, 0xFF},
	"gray":    {0x5F, 0x63, 0x68, 0xFF},
}

// parseColor accepts "#rgb", "#rrggbb", "#rrggbbaa" or one of the names above.
func parseColor(s string) (color.NRGBA, error) {
	t := strings.ToLower(strings.TrimSpace(s))
	if t == "" {
		return namedColors["red"], nil
	}
	if c, ok := namedColors[t]; ok {
		return c, nil
	}
	if !strings.HasPrefix(t, "#") {
		return color.NRGBA{}, fmt.Errorf("unknown colour %q (use #rrggbb or one of red, orange, yellow, green, blue, purple, magenta, black, white, grey)", s)
	}
	hexs := t[1:]
	switch len(hexs) {
	case 3:
		hexs = string([]byte{hexs[0], hexs[0], hexs[1], hexs[1], hexs[2], hexs[2], 'f', 'f'})
	case 6:
		hexs += "ff"
	case 8:
	default:
		return color.NRGBA{}, fmt.Errorf("colour %q must have 3, 6 or 8 hex digits", s)
	}
	v, err := strconv.ParseUint(hexs, 16, 64)
	if err != nil {
		return color.NRGBA{}, fmt.Errorf("colour %q is not hexadecimal", s)
	}
	return color.NRGBA{
		R: uint8(v >> 24), G: uint8(v >> 16), B: uint8(v >> 8), A: uint8(v),
	}, nil
}

// darken returns c scaled towards black by f (0 = unchanged, 1 = black).
func darken(c color.NRGBA, f float64) color.NRGBA {
	k := 1 - f
	return color.NRGBA{
		R: uint8(float64(c.R) * k),
		G: uint8(float64(c.G) * k),
		B: uint8(float64(c.B) * k),
		A: c.A,
	}
}

// contrastInk returns black or white, whichever is readable on c.
func contrastInk(c color.NRGBA) color.NRGBA {
	// Rec. 601 luma.
	y := 0.299*float64(c.R) + 0.587*float64(c.G) + 0.114*float64(c.B)
	if y > 150 {
		return color.NRGBA{0x11, 0x11, 0x11, 0xFF}
	}
	return color.NRGBA{0xFF, 0xFF, 0xFF, 0xFF}
}

// ---------------------------------------------------------------------------
// PNG helpers
// ---------------------------------------------------------------------------

// decodePNGFile decodes path as PNG and also returns the SHA-256 of the file
// bytes. The file is opened read-only and is never written.
func decodePNGFile(path string) (image.Image, string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, "", err
	}
	img, err := png.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, "", err
	}
	sum := sha256.Sum256(data)
	return img, hex.EncodeToString(sum[:]), nil
}

func encodePNG(img image.Image) ([]byte, error) {
	var buf bytes.Buffer
	enc := png.Encoder{CompressionLevel: png.DefaultCompression}
	if err := enc.Encode(&buf, img); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// toRGBA copies src into a fresh RGBA image anchored at (0,0).
// The source image is never mutated.
func toRGBA(src image.Image) *image.RGBA {
	b := src.Bounds()
	dst := image.NewRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
	draw.Draw(dst, dst.Bounds(), src, b.Min, draw.Src)
	return dst
}

// ---------------------------------------------------------------------------
// Rasterisers
//
// Everything below composites with a coverage value in [0,1]; that is where the
// anti-aliasing comes from. No third-party rasteriser is involved.
// ---------------------------------------------------------------------------

// blend paints c over the pixel at (x,y) with coverage a in [0,1].
func blend(dst *image.RGBA, x, y int, c color.NRGBA, a float64) {
	if a <= 0 {
		return
	}
	if a > 1 {
		a = 1
	}
	if !(image.Point{x, y}.In(dst.Bounds())) {
		return
	}
	alpha := a * float64(c.A) / 255
	i := dst.PixOffset(x, y)
	// dst is opaque in practice, but composite correctly regardless.
	dr := float64(dst.Pix[i+0])
	dg := float64(dst.Pix[i+1])
	db := float64(dst.Pix[i+2])
	da := float64(dst.Pix[i+3])
	dst.Pix[i+0] = uint8(float64(c.R)*alpha + dr*(1-alpha) + 0.5)
	dst.Pix[i+1] = uint8(float64(c.G)*alpha + dg*(1-alpha) + 0.5)
	dst.Pix[i+2] = uint8(float64(c.B)*alpha + db*(1-alpha) + 0.5)
	dst.Pix[i+3] = uint8(255*alpha + da*(1-alpha) + 0.5)
}

// coverage maps a signed distance (negative inside) to an alpha with a
// one-pixel-wide linear ramp centred on the edge.
func coverage(dist float64) float64 {
	a := 0.5 - dist
	if a < 0 {
		return 0
	}
	if a > 1 {
		return 1
	}
	return a
}

// fillCircle draws a filled, anti-aliased disc of radius r at (cx,cy).
func fillCircle(dst *image.RGBA, cx, cy, r float64, c color.NRGBA) {
	x0 := int(math.Floor(cx - r - 1))
	x1 := int(math.Ceil(cx + r + 1))
	y0 := int(math.Floor(cy - r - 1))
	y1 := int(math.Ceil(cy + r + 1))
	for y := y0; y <= y1; y++ {
		for x := x0; x <= x1; x++ {
			dx := float64(x) + 0.5 - cx
			dy := float64(y) + 0.5 - cy
			d := math.Hypot(dx, dy) - r
			blend(dst, x, y, c, coverage(d))
		}
	}
}

// strokeCircle draws an anti-aliased ring of the given width centred on radius r.
func strokeCircle(dst *image.RGBA, cx, cy, r, width float64, c color.NRGBA) {
	half := width / 2
	x0 := int(math.Floor(cx - r - half - 1))
	x1 := int(math.Ceil(cx + r + half + 1))
	y0 := int(math.Floor(cy - r - half - 1))
	y1 := int(math.Ceil(cy + r + half + 1))
	for y := y0; y <= y1; y++ {
		for x := x0; x <= x1; x++ {
			dx := float64(x) + 0.5 - cx
			dy := float64(y) + 0.5 - cy
			d := math.Abs(math.Hypot(dx, dy)-r) - half
			blend(dst, x, y, c, coverage(d))
		}
	}
}

// strokeLine draws an anti-aliased line of the given width with round caps.
func strokeLine(dst *image.RGBA, x0, y0, x1, y1, width float64, c color.NRGBA) {
	half := width / 2
	minX := int(math.Floor(math.Min(x0, x1) - half - 1))
	maxX := int(math.Ceil(math.Max(x0, x1) + half + 1))
	minY := int(math.Floor(math.Min(y0, y1) - half - 1))
	maxY := int(math.Ceil(math.Max(y0, y1) + half + 1))
	for y := minY; y <= maxY; y++ {
		for x := minX; x <= maxX; x++ {
			d := distToSegment(float64(x)+0.5, float64(y)+0.5, x0, y0, x1, y1) - half
			blend(dst, x, y, c, coverage(d))
		}
	}
}

func distToSegment(px, py, x0, y0, x1, y1 float64) float64 {
	vx, vy := x1-x0, y1-y0
	wx, wy := px-x0, py-y0
	den := vx*vx + vy*vy
	if den == 0 {
		return math.Hypot(wx, wy)
	}
	t := (wx*vx + wy*vy) / den
	if t < 0 {
		t = 0
	} else if t > 1 {
		t = 1
	}
	return math.Hypot(px-(x0+t*vx), py-(y0+t*vy))
}

// strokeRect draws an anti-aliased rectangle outline. The stroke is centred on
// the rectangle's edges.
func strokeRect(dst *image.RGBA, x, y, w, h, width float64, c color.NRGBA) {
	half := width / 2
	minX := int(math.Floor(x - half - 1))
	maxX := int(math.Ceil(x + w + half + 1))
	minY := int(math.Floor(y - half - 1))
	maxY := int(math.Ceil(y + h + half + 1))
	for py := minY; py <= maxY; py++ {
		for px := minX; px <= maxX; px++ {
			fx := float64(px) + 0.5
			fy := float64(py) + 0.5
			// Signed distance to the rectangle border (0 on the border).
			dx := math.Max(x-fx, fx-(x+w))
			dy := math.Max(y-fy, fy-(y+h))
			var d float64
			if dx <= 0 && dy <= 0 {
				d = -math.Max(dx, dy) // inside: distance to nearest edge
			} else {
				d = math.Hypot(math.Max(dx, 0), math.Max(dy, 0))
			}
			blend(dst, px, py, c, coverage(d-half))
		}
	}
}

// fillTriangle draws an anti-aliased filled triangle by 4x4 supersampling of the
// three edge half-plane tests.
func fillTriangle(dst *image.RGBA, ax, ay, bx, by, cx, cy float64, col color.NRGBA) {
	minX := int(math.Floor(math.Min(ax, math.Min(bx, cx)) - 1))
	maxX := int(math.Ceil(math.Max(ax, math.Max(bx, cx)) + 1))
	minY := int(math.Floor(math.Min(ay, math.Min(by, cy)) - 1))
	maxY := int(math.Ceil(math.Max(ay, math.Max(by, cy)) + 1))
	area := edge(ax, ay, bx, by, cx, cy)
	if area == 0 {
		return
	}
	sign := 1.0
	if area < 0 {
		sign = -1
	}
	const ss = 4
	for y := minY; y <= maxY; y++ {
		for x := minX; x <= maxX; x++ {
			hits := 0
			for sy := 0; sy < ss; sy++ {
				for sx := 0; sx < ss; sx++ {
					px := float64(x) + (float64(sx)+0.5)/ss
					py := float64(y) + (float64(sy)+0.5)/ss
					e0 := sign * edge(ax, ay, bx, by, px, py)
					e1 := sign * edge(bx, by, cx, cy, px, py)
					e2 := sign * edge(cx, cy, ax, ay, px, py)
					if e0 >= 0 && e1 >= 0 && e2 >= 0 {
						hits++
					}
				}
			}
			if hits > 0 {
				blend(dst, x, y, col, float64(hits)/(ss*ss))
			}
		}
	}
}

func edge(ax, ay, bx, by, px, py float64) float64 {
	return (bx-ax)*(py-ay) - (by-ay)*(px-ax)
}

// drawText draws s with the embedded bitmap glyphs, scaled by an integer
// factor, with its top-left at (x,y). It returns the number of unknown runes,
// which the caller reports as a warning.
func drawText(dst *image.RGBA, s string, x, y, scale int, c color.NRGBA) int {
	if scale < 1 {
		scale = 1
	}
	unknown := 0
	penX := x
	for _, r := range s {
		if _, ok := glyphFor(r); !ok {
			unknown++
		}
		for gy := 0; gy < glyphH; gy++ {
			for gx := 0; gx < glyphW; gx++ {
				if !glyphPixels(r, gx, gy) {
					continue
				}
				for sy := 0; sy < scale; sy++ {
					for sx := 0; sx < scale; sx++ {
						blend(dst, penX+gx*scale+sx, y+gy*scale+sy, c, 1)
					}
				}
			}
		}
		penX += (glyphW + 1) * scale
	}
	return unknown
}

// ---------------------------------------------------------------------------
// Callout painting
// ---------------------------------------------------------------------------

// resolvedCallout is a callout with defaults filled in, ready to draw and to
// list in the document's legend.
type resolvedCallout struct {
	Callout
	Number int // sequential circle number, 0 for rect/arrow
	Label  string
	Color  color.NRGBA
}

// resolveCallouts fills in default radius, colour and auto-numbering.
func resolveCallouts(cs []Callout) []resolvedCallout {
	out := make([]resolvedCallout, 0, len(cs))
	n := 0
	for _, c := range cs {
		rc := resolvedCallout{Callout: c}
		rc.Color, _ = parseColor(c.Color)
		if c.Kind == "circle" {
			n++
			rc.Number = n
			if strings.TrimSpace(c.Label) == "" {
				rc.Label = strconv.Itoa(n)
			} else {
				rc.Label = c.Label
			}
			if rc.R == 0 {
				rc.R = defaultCircleRadius
			}
		} else {
			rc.Label = strings.TrimSpace(c.Label)
		}
		out = append(out, rc)
	}
	return out
}

// annotate returns a NEW image with the callouts drawn on it. src is only read.
func annotate(src image.Image, cs []resolvedCallout) (*image.RGBA, int) {
	dst := toRGBA(src)
	unknown := 0
	for _, c := range cs {
		switch c.Kind {
		case "circle":
			unknown += drawCircleCallout(dst, c)
		case "rect":
			drawRectCallout(dst, c)
		case "arrow":
			unknown += drawArrowCallout(dst, c)
		}
	}
	return dst, unknown
}

func drawCircleCallout(dst *image.RGBA, c resolvedCallout) int {
	cx := float64(c.X) + 0.5
	cy := float64(c.Y) + 0.5
	r := float64(c.R)
	// White halo, then the filled disc, then a darker ring: readable on any
	// background.
	fillCircle(dst, cx, cy, r+2, color.NRGBA{0xFF, 0xFF, 0xFF, 0xFF})
	fillCircle(dst, cx, cy, r, c.Color)
	strokeCircle(dst, cx, cy, r, 2, darken(c.Color, 0.35))

	ink := contrastInk(c.Color)
	scale := int(math.Max(1, math.Floor((r*1.15)/glyphH)))
	for scale > 1 && float64(textWidth(c.Label)*scale) > 1.7*r {
		scale--
	}
	tw := textWidth(c.Label) * scale
	th := glyphH * scale
	tx := int(math.Round(cx)) - tw/2
	ty := int(math.Round(cy)) - th/2
	return drawText(dst, c.Label, tx, ty, scale, ink)
}

func drawRectCallout(dst *image.RGBA, c resolvedCallout) {
	x, y := float64(c.X), float64(c.Y)
	w, h := float64(c.W), float64(c.H)
	// Halo first so the outline stays visible on a busy screenshot.
	strokeRect(dst, x, y, w, h, strokeWidth+3, color.NRGBA{0xFF, 0xFF, 0xFF, 0xC0})
	strokeRect(dst, x, y, w, h, strokeWidth, c.Color)
}

func drawArrowCallout(dst *image.RGBA, c resolvedCallout) int {
	x0, y0 := float64(c.X)+0.5, float64(c.Y)+0.5
	x1, y1 := float64(c.X2)+0.5, float64(c.Y2)+0.5
	dx, dy := x1-x0, y1-y0
	length := math.Hypot(dx, dy)
	if length == 0 {
		return 0
	}
	ux, uy := dx/length, dy/length

	headLen := math.Min(18, length*0.45)
	headHalf := headLen * 0.5
	// Shaft stops where the head begins so the join stays crisp.
	sx := x1 - ux*headLen
	sy := y1 - uy*headLen
	// Perpendicular unit vector.
	px, py := -uy, ux

	halo := color.NRGBA{0xFF, 0xFF, 0xFF, 0xC0}
	strokeLine(dst, x0, y0, sx, sy, strokeWidth+3, halo)
	fillTriangle(dst,
		x1+ux*1.5, y1+uy*1.5,
		sx+px*(headHalf+1.5), sy+py*(headHalf+1.5),
		sx-px*(headHalf+1.5), sy-py*(headHalf+1.5), halo)

	strokeLine(dst, x0, y0, sx, sy, strokeWidth, c.Color)
	fillTriangle(dst, x1, y1,
		sx+px*headHalf, sy+py*headHalf,
		sx-px*headHalf, sy-py*headHalf, c.Color)

	if c.Label == "" {
		return 0
	}
	scale := 2
	tw := textWidth(c.Label) * scale
	th := glyphH * scale
	// Label sits behind the tail, clear of the shaft.
	lx := int(math.Round(x0 - ux*float64(th) - float64(tw)/2))
	ly := int(math.Round(y0 - uy*float64(th) - float64(th)/2))
	pad := 3
	fillRectSolid(dst, lx-pad, ly-pad, tw+2*pad, th+2*pad, color.NRGBA{0xFF, 0xFF, 0xFF, 0xE6})
	return drawText(dst, c.Label, lx, ly, scale, darken(c.Color, 0.15))
}

func fillRectSolid(dst *image.RGBA, x, y, w, h int, c color.NRGBA) {
	for py := y; py < y+h; py++ {
		for px := x; px < x+w; px++ {
			blend(dst, px, py, c, 1)
		}
	}
}