package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
)

// rect is a stamp rectangle in PDF points, bottom-left origin.
type rect struct {
	X float64 `json:"x"`
	Y float64 `json:"y"`
	W float64 `json:"w"`
	H float64 `json:"h"`
}

// stamp is one stamp placed on one or more pages.
type stamp struct {
	ID      string  `json:"id"`
	Type    string  `json:"type"`              // "text" or "image"
	Content string  `json:"content,omitempty"` // text stamps
	Src     string  `json:"src,omitempty"`     // image stamps: PNG or JPEG path
	Pages   []int   `json:"pages"`             // 1-based page numbers
	Rect    rect    `json:"rect"`
	Opacity float64 `json:"opacity"`
}

// plan is the on-disk stamp plan.
type plan struct {
	Version int     `json:"version"`
	Input   string  `json:"input"`
	Stamps  []stamp `json:"stamps"`
}

const (
	planVersion  = 2
	typeText     = "text"
	typeImage    = "image"
	maxPlanBytes = 8 << 20
	maxStamps    = 1_000
	maxPageRefs  = 100_000
	maxTextBytes = 64 << 10
)

// newPlan returns an empty plan for the given input PDF.
func newPlan(input string) *plan {
	return &plan{
		Version: planVersion,
		Input:   input,
		Stamps:  []stamp{},
	}
}

// loadPlan reads and validates a plan from path.
func loadPlan(path string) (*plan, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return readPlan(f, path)
}

func readPlan(r io.Reader, name string) (*plan, error) {
	b, err := io.ReadAll(io.LimitReader(r, maxPlanBytes+1))
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", name, err)
	}
	return parsePlan(b, name)
}

// parsePlan strictly decodes and validates plan JSON; name is used in errors.
func parsePlan(b []byte, name string) (*plan, error) {
	if len(b) > maxPlanBytes {
		return nil, fmt.Errorf("plan %s is larger than %d bytes", name, maxPlanBytes)
	}
	var p plan
	dec := json.NewDecoder(bytes.NewReader(b))
	dec.DisallowUnknownFields()
	if err := dec.Decode(&p); err != nil {
		return nil, fmt.Errorf("parse %s: %w", name, err)
	}
	if err := dec.Decode(&struct{}{}); err != io.EOF {
		if err == nil {
			return nil, fmt.Errorf("parse %s: multiple JSON values", name)
		}
		return nil, fmt.Errorf("parse %s: trailing data: %w", name, err)
	}
	if err := p.validate(); err != nil {
		return nil, fmt.Errorf("invalid plan %s: %w", name, err)
	}
	return &p, nil
}

// save writes the plan to path as indented JSON.
func (p *plan) save(path string) error {
	b, err := p.marshal()
	if err != nil {
		return err
	}
	return writeFileAtomic(path, b, 0o644)
}

func (p *plan) marshal() ([]byte, error) {
	if err := p.validate(); err != nil {
		return nil, err
	}
	b, err := json.MarshalIndent(p, "", "  ")
	if err != nil {
		return nil, err
	}
	b = append(b, '\n')
	if len(b) > maxPlanBytes {
		return nil, fmt.Errorf("encoded plan is larger than %d bytes", maxPlanBytes)
	}
	return b, nil
}

// writeFileAtomic writes b to path via a temp file in the same directory,
// where rename is atomic on local filesystems. The temp file is synced before
// the rename, so an interrupted write leaves it behind rather than a truncated
// file at path.
func writeFileAtomic(path string, b []byte, perm os.FileMode) error {
	dir := filepath.Dir(path)
	f, err := os.CreateTemp(dir, "."+filepath.Base(path)+".tmp-*")
	if err != nil {
		return err
	}
	tmp := f.Name()
	defer os.Remove(tmp) // no-op after a successful rename
	if _, err := f.Write(b); err != nil {
		f.Close()
		return err
	}
	if err := f.Chmod(perm); err != nil {
		f.Close()
		return err
	}
	if err := f.Sync(); err != nil {
		f.Close()
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		return err
	}
	return syncDir(dir)
}

func syncDir(path string) error {
	d, err := os.Open(path)
	if err != nil {
		return err
	}
	if err := d.Sync(); err != nil {
		d.Close()
		return err
	}
	return d.Close()
}

// validate checks the plan's invariants.
func (p *plan) validate() error {
	if p.Version != planVersion {
		return fmt.Errorf("version %d unsupported (want %d)", p.Version, planVersion)
	}
	if p.Input == "" {
		return fmt.Errorf("input is empty")
	}
	if len(p.Stamps) > maxStamps {
		return fmt.Errorf("too many stamps (%d; limit is %d)", len(p.Stamps), maxStamps)
	}
	pageRefs := 0
	for i := range p.Stamps {
		pageRefs += len(p.Stamps[i].Pages)
		if pageRefs > maxPageRefs {
			return fmt.Errorf("too many page references (%d; limit is %d)", pageRefs, maxPageRefs)
		}
		if err := p.Stamps[i].validate(); err != nil {
			return fmt.Errorf("stamp %d (%s): %w", i, p.Stamps[i].ID, err)
		}
	}
	return nil
}

func (s *stamp) validate() error {
	switch s.Type {
	case typeText:
		if s.Content == "" {
			return fmt.Errorf("text stamp has empty content")
		}
		if len(s.Content) > maxTextBytes {
			return fmt.Errorf("text is larger than %d bytes", maxTextBytes)
		}
		// Text is written as raw bytes into a WinAnsi-encoded base-14 string,
		// so anything outside printable ASCII would render as wrong glyphs.
		for _, r := range s.Content {
			if r < 0x20 || r > 0x7e {
				return fmt.Errorf("text contains %q; only printable ASCII is supported (no font is embedded)", r)
			}
		}
	case typeImage:
		if s.Src == "" {
			return fmt.Errorf("image stamp has empty src")
		}
	default:
		return fmt.Errorf("unknown type %q", s.Type)
	}
	if len(s.Pages) == 0 {
		return fmt.Errorf("no pages")
	}
	seen := make(map[int]struct{}, len(s.Pages))
	for _, pg := range s.Pages {
		if pg < 1 {
			return fmt.Errorf("page %d out of range (1-based)", pg)
		}
		if _, ok := seen[pg]; ok {
			return fmt.Errorf("duplicate page %d", pg)
		}
		seen[pg] = struct{}{}
	}
	for _, v := range []float64{s.Rect.X, s.Rect.Y, s.Rect.W, s.Rect.H, s.Opacity} {
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return fmt.Errorf("rect and opacity must be finite (got %g)", v)
		}
	}
	if s.Rect.W <= 0 || s.Rect.H <= 0 {
		return fmt.Errorf("rect must have positive w and h")
	}
	if s.Opacity < 0 || s.Opacity > 1 {
		return fmt.Errorf("opacity %g out of range [0,1]", s.Opacity)
	}
	return nil
}

// toAbsolute rewrites Input and image srcs to absolute paths. Called before
// saving so a plan applies correctly no matter which directory it is run from.
func (p *plan) toAbsolute() {
	p.Input = absPath(p.Input)
	for i := range p.Stamps {
		if p.Stamps[i].Type == typeImage {
			p.Stamps[i].Src = absPath(p.Stamps[i].Src)
		}
	}
}

// resolveRelativeTo makes relative Input and image srcs relative to base (the
// plan file's directory), so hand-written plans with relative paths apply
// correctly regardless of the current working directory.
func (p *plan) resolveRelativeTo(base string) {
	p.Input = joinIfRel(base, p.Input)
	for i := range p.Stamps {
		if p.Stamps[i].Type == typeImage {
			p.Stamps[i].Src = joinIfRel(base, p.Stamps[i].Src)
		}
	}
}

func absPath(path string) string {
	if path == "" {
		return path
	}
	if a, err := filepath.Abs(path); err == nil {
		return a
	}
	return path
}

func joinIfRel(base, path string) string {
	if path == "" || filepath.IsAbs(path) {
		return path
	}
	return filepath.Join(base, path)
}

// maxPage returns the highest 1-based page referenced by any stamp.
func (p *plan) maxPage() int {
	mx := 0
	for _, s := range p.Stamps {
		for _, pg := range s.Pages {
			if pg > mx {
				mx = pg
			}
		}
	}
	return mx
}