package main

import (
	"bytes"
	"compress/zlib"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"math"
	"os"
	"strconv"
	"strings"

	"github.com/pdfcpu/pdfcpu/pkg/pdfcpu/types"
)

// Image limits bound the decoded working set, not the compressed file size.
const (
	maxImagePixels      = 50_000_000
	maxTotalImagePixels = 100_000_000
)

// buildOverlay constructs an in-memory PDF with one page per source page,
// each sized to the source page and drawing that page's stamps at absolute
// coordinates (bottom-left origin, points). It is meant to be overlaid 1:1
// onto the source by pdfcpu (multi-stamp, pos:c, sc:1 abs).
func buildOverlay(p *plan, dims []types.Dim) ([]byte, error) {
	n := len(dims)
	byPage := map[int][]*stamp{}
	haveText := false
	for i := range p.Stamps {
		s := &p.Stamps[i]
		if s.Type == typeText {
			haveText = true
		}
		for _, pg := range s.Pages {
			if pg >= 1 && pg <= n {
				byPage[pg] = append(byPage[pg], s)
			}
		}
	}

	w := &pdfWriter{}
	catalog := w.reserve()
	pages := w.reserve()

	// Shared font (Helvetica base-14, no embedding).
	fontNum := 0
	if haveText {
		fontNum = w.add([]byte("<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica /Encoding /WinAnsiEncoding >>"))
	}

	// One ExtGState per distinct opacity, created before the shared Resources
	// dict so that dict lists every ExtGState the pages reference.
	gs := map[string]string{} // opacity key -> ExtGState name
	var gsEntries []string
	for i := range p.Stamps {
		k := opacityKey(p.Stamps[i].Opacity)
		if _, ok := gs[k]; ok {
			continue
		}
		name := fmt.Sprintf("GS%d", len(gs))
		num := w.add([]byte(fmt.Sprintf("<< /Type /ExtGState /ca %s /CA %s >>", k, k)))
		gs[k] = name
		gsEntries = append(gsEntries, fmt.Sprintf("/%s %d 0 R", name, num))
	}

	// One image XObject per distinct src.
	imgs := map[string]imageRes{}
	var xobjEntries []string
	remainingPixels := int64(maxTotalImagePixels)
	for i := range p.Stamps {
		s := &p.Stamps[i]
		if s.Type != typeImage {
			continue
		}
		if _, ok := imgs[s.Src]; ok {
			continue
		}
		res, pixels, err := embedImage(w, s.Src, remainingPixels)
		if err != nil {
			return nil, fmt.Errorf("image stamp %q: %w", s.Src, err)
		}
		remainingPixels -= pixels
		res.name = fmt.Sprintf("Im%d", len(imgs))
		imgs[s.Src] = res
		xobjEntries = append(xobjEntries, fmt.Sprintf("/%s %d 0 R", res.name, res.num))
	}

	var rb strings.Builder
	rb.WriteString("<< /ProcSet [/PDF /Text /ImageB /ImageC /ImageI]")
	if fontNum != 0 {
		fmt.Fprintf(&rb, " /Font << /F1 %d 0 R >>", fontNum)
	}
	if len(gsEntries) > 0 {
		fmt.Fprintf(&rb, " /ExtGState << %s >>", strings.Join(gsEntries, " "))
	}
	if len(xobjEntries) > 0 {
		fmt.Fprintf(&rb, " /XObject << %s >>", strings.Join(xobjEntries, " "))
	}
	rb.WriteString(" >>")
	resNum := w.add([]byte(rb.String()))

	pageNums := make([]int, n)
	for i := 0; i < n; i++ {
		content := pageContent(byPage[i+1], gs, imgs)
		cNum := w.add(makeStream("", content))
		dict := fmt.Sprintf(
			"<< /Type /Page /Parent %d 0 R /MediaBox [0 0 %s %s] /Resources %d 0 R /Contents %d 0 R >>",
			pages, ff(dims[i].Width), ff(dims[i].Height), resNum, cNum)
		pageNums[i] = w.add([]byte(dict))
	}

	kids := make([]string, n)
	for i, num := range pageNums {
		kids[i] = fmt.Sprintf("%d 0 R", num)
	}
	w.set(pages, []byte(fmt.Sprintf("<< /Type /Pages /Kids [%s] /Count %d >>", strings.Join(kids, " "), n)))
	w.set(catalog, []byte(fmt.Sprintf("<< /Type /Catalog /Pages %d 0 R >>", pages)))

	return w.bytes(), nil
}

// pageContent renders all stamps for one page into a content stream.
func pageContent(stamps []*stamp, gs map[string]string, imgs map[string]imageRes) []byte {
	var b bytes.Buffer
	for _, s := range stamps {
		state := gs[opacityKey(s.Opacity)]
		switch s.Type {
		case typeText:
			size, dx, dy := fitText(s.Content, s.Rect.W, s.Rect.H)
			// cm: translate to the rect origin; Td centers the line in the box.
			fmt.Fprintf(&b, "q /%s gs 1 0 0 1 %s %s cm BT /F1 %s Tf 0 g %s %s Td (%s) Tj ET Q\n",
				state, ff(s.Rect.X), ff(s.Rect.Y),
				ff(size), ff(dx), ff(dy), textEscaper.Replace(s.Content))
		case typeImage:
			res := imgs[s.Src]
			sc := math.Min(s.Rect.W/float64(res.w), s.Rect.H/float64(res.h))
			fmt.Fprintf(&b, "q /%s gs %s 0 0 %s %s %s cm /%s Do Q\n",
				state, ff(float64(res.w)*sc), ff(float64(res.h)*sc),
				ff(s.Rect.X), ff(s.Rect.Y), res.name)
		}
	}
	return b.Bytes()
}

// opacityKey formats an opacity for the ExtGState dict; stamps share a state
// when they format identically.
func opacityKey(v float64) string { return strconv.FormatFloat(v, 'f', 3, 64) }

// helvWidths holds Helvetica glyph widths for printable ASCII (0x20..0x7e)
// in 1/1000 em, from the Adobe core-14 AFM metrics. Validation restricts
// text content to this range, so exact fitting needs no embedded font.
var helvWidths = [95]int{
	278, 278, 355, 556, 556, 889, 667, 191, 333, 333, 389, 584, 278, 333, 278, 278,
	556, 556, 556, 556, 556, 556, 556, 556, 556, 556, 278, 278, 584, 584, 584, 556,
	1015, 667, 667, 722, 722, 667, 611, 778, 722, 278, 500, 667, 556, 833, 722, 778,
	667, 778, 722, 667, 611, 722, 667, 944, 667, 667, 611, 278, 278, 278, 469, 556,
	333, 556, 556, 500, 556, 556, 278, 556, 556, 222, 222, 500, 222, 833, 556, 556,
	556, 556, 333, 500, 278, 556, 500, 722, 500, 500, 500, 334, 260, 334, 584,
}

// helvCapHeight is Helvetica's cap height in em units (AFM CapHeight 718).
const helvCapHeight = 0.718

// textWidth returns the width of s in em units (multiply by the font size).
func textWidth(s string) float64 {
	w := 0
	for i := 0; i < len(s); i++ {
		w += helvWidths[s[i]-0x20]
	}
	return float64(w) / 1000
}

// fitText sizes text to its rectangle using the real metrics: height drives
// the size, shrunk if the line would overflow the width. dx/dy center the
// line inside the box (offsets from the rect origin).
func fitText(text string, w, h float64) (size, dx, dy float64) {
	uw := textWidth(text)
	size = 0.72 * h
	if uw*size > w {
		size = w / uw
	}
	return size, (w - uw*size) / 2, (h - helvCapHeight*size) / 2
}

// textEscaper escapes PDF string delimiters; validation already restricts
// content to printable ASCII.
var textEscaper = strings.NewReplacer("\\", "\\\\", "(", "\\(", ")", "\\)")

// ff formats a float for PDF content with trailing zeros trimmed.
func ff(v float64) string {
	s := strconv.FormatFloat(v, 'f', 4, 64)
	if strings.Contains(s, ".") {
		s = strings.TrimRight(strings.TrimRight(s, "0"), ".")
	}
	if s == "-0" {
		s = "0"
	}
	return s
}

type imageRes struct {
	name string // XObject name in the resources dict
	num  int    // image XObject number
	w, h int    // pixel dimensions
}

// embedImage decodes a PNG/JPEG and writes an RGB image XObject (plus an
// 8-bit DeviceGray SMask when the image has transparency). Header dimensions
// are checked before decoding so a huge file fails fast, not out of memory.
func embedImage(w *pdfWriter, path string, remainingPixels int64) (imageRes, int64, error) {
	f, err := os.Open(path)
	if err != nil {
		return imageRes{}, 0, err
	}
	defer f.Close()
	cfg, _, err := decodeImageConfig(f)
	if err != nil {
		return imageRes{}, 0, err
	}
	pixels := int64(cfg.Width) * int64(cfg.Height)
	if pixels > remainingPixels {
		return imageRes{}, 0, fmt.Errorf("images exceed the total limit of %d pixels", maxTotalImagePixels)
	}
	if _, err := f.Seek(0, io.SeekStart); err != nil {
		return imageRes{}, 0, err
	}
	img, format, err := image.Decode(f)
	if err != nil {
		return imageRes{}, 0, err
	}
	if !supportedImageFormat(format) {
		return imageRes{}, 0, fmt.Errorf("image format %q is not supported (want PNG or JPEG)", format)
	}
	bnds := img.Bounds()
	pw, ph := bnds.Dx(), bnds.Dy()
	if pw != cfg.Width || ph != cfg.Height {
		return imageRes{}, 0, fmt.Errorf("decoded dimensions %dx%d do not match header %dx%d", pw, ph, cfg.Width, cfg.Height)
	}
	rgb := make([]byte, 0, 3*pw*ph)
	hasAlpha := true
	if opaque, ok := img.(interface{ Opaque() bool }); ok {
		hasAlpha = !opaque.Opaque()
	}
	var alpha []byte
	if hasAlpha {
		alpha = make([]byte, 0, pw*ph)
	}
	for y := bnds.Min.Y; y < bnds.Max.Y; y++ {
		for x := bnds.Min.X; x < bnds.Max.X; x++ {
			r, g, b, a := img.At(x, y).RGBA() // 16-bit, alpha-premultiplied
			a8 := byte(a >> 8)
			// Un-premultiply so the SMask carries the transparency.
			if a == 0 {
				rgb = append(rgb, 0, 0, 0)
			} else {
				rgb = append(rgb, byte((r*0xffff/a)>>8), byte((g*0xffff/a)>>8), byte((b*0xffff/a)>>8))
			}
			if hasAlpha {
				alpha = append(alpha, a8)
			}
		}
	}

	smask := 0
	if hasAlpha {
		smask = w.add(makeStream(
			fmt.Sprintf("/Type /XObject /Subtype /Image /Width %d /Height %d /ColorSpace /DeviceGray /BitsPerComponent 8 /Filter /FlateDecode", pw, ph),
			flate(alpha)))
	}
	dict := fmt.Sprintf("/Type /XObject /Subtype /Image /Width %d /Height %d /ColorSpace /DeviceRGB /BitsPerComponent 8 /Filter /FlateDecode", pw, ph)
	if smask != 0 {
		dict += fmt.Sprintf(" /SMask %d 0 R", smask)
	}
	num := w.add(makeStream(dict, flate(rgb)))
	return imageRes{num: num, w: pw, h: ph}, pixels, nil
}

// decodeImageConfig reads an image header and enforces maxImagePixels. It is
// shared by embedImage and the picker's pick-time check so both agree.
func decodeImageConfig(f *os.File) (image.Config, string, error) {
	cfg, format, err := image.DecodeConfig(f)
	if err != nil {
		return cfg, "", err
	}
	if !supportedImageFormat(format) {
		return cfg, format, fmt.Errorf("image format %q is not supported (want PNG or JPEG)", format)
	}
	if cfg.Width < 1 || cfg.Height < 1 || int64(cfg.Width) > int64(maxImagePixels)/int64(cfg.Height) {
		return cfg, format, fmt.Errorf("image is %dx%d; the limit is %d pixels", cfg.Width, cfg.Height, maxImagePixels)
	}
	return cfg, format, nil
}

func supportedImageFormat(format string) bool {
	return format == "png" || format == "jpeg"
}

func flate(b []byte) []byte {
	var buf bytes.Buffer
	zw := zlib.NewWriter(&buf)
	zw.Write(b)
	zw.Close()
	return buf.Bytes()
}

// --- minimal PDF object writer ---

type pdfWriter struct {
	objs [][]byte // objs[i] is the body of object number i+1
}

func (w *pdfWriter) reserve() int { w.objs = append(w.objs, nil); return len(w.objs) }
func (w *pdfWriter) add(body []byte) int {
	w.objs = append(w.objs, body)
	return len(w.objs)
}
func (w *pdfWriter) set(num int, body []byte) { w.objs[num-1] = body }

// makeStream builds a stream object body: dict (without /Length) + stream data.
func makeStream(dictInner string, data []byte) []byte {
	var b bytes.Buffer
	if dictInner == "" {
		fmt.Fprintf(&b, "<< /Length %d >>\nstream\n", len(data))
	} else {
		fmt.Fprintf(&b, "<< %s /Length %d >>\nstream\n", dictInner, len(data))
	}
	b.Write(data)
	b.WriteString("\nendstream")
	return b.Bytes()
}

// bytes serializes all objects with a classic cross-reference table.
func (w *pdfWriter) bytes() []byte {
	var b bytes.Buffer
	b.WriteString("%PDF-1.7\n%\xe2\xe3\xcf\xd3\n")
	offsets := make([]int, len(w.objs))
	for i, body := range w.objs {
		offsets[i] = b.Len()
		fmt.Fprintf(&b, "%d 0 obj\n", i+1)
		b.Write(body)
		b.WriteString("\nendobj\n")
	}
	xref := b.Len()
	fmt.Fprintf(&b, "xref\n0 %d\n", len(w.objs)+1)
	b.WriteString("0000000000 65535 f \n")
	for _, off := range offsets {
		fmt.Fprintf(&b, "%010d 00000 n \n", off)
	}
	fmt.Fprintf(&b, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(w.objs)+1, xref)
	return b.Bytes()
}