package ui

import (
	"bytes"
	"encoding/base64"
	"fmt"
	xdraw "golang.org/x/image/draw"
	"image"
	"image/png"
	"strings"
	"time"

	"github.com/gdamore/tcell/v2"
	"github.com/rivo/tview"
)

// The Kitty graphics protocol with Unicode placeholders, as supported by
// kitty and Ghostty: the poster is uploaded to the terminal once, and then
// shown by writing special placeholder characters into ordinary text cells.
// Because the picture is made of text cells it behaves like text: popups and
// dialogs drawn over it cover it, it vanishes when the view changes, and the
// normal screen diffing keeps it in place. The poster is shown at the
// terminal's full resolution.

// placeholderRune marks a cell that shows a piece of an image. It is followed
// by combining characters that give the row and column in the image; the
// image id is the cell's foreground color.
const placeholderRune = '\U0010EEEE'

const kittyChunk = 4096 // bytes of base64 per escape sequence, the protocol maximum

// kittyGraphics uploads images to the terminal and keeps track of them.
//
// Every footprint of a poster (the part of it that is shown and the number of
// cells it covers) is uploaded as its own image with its own id and placed once.
// An existing placement is never changed: terminals do not reliably replace the
// size of a placement that is already there (Ghostty kept showing the old, small
// one after the window was enlarged), whereas a new image id always starts clean.
type kittyGraphics struct {
	write  func([]byte) error // writes to the terminal, in step with the screen
	nextID uint32
	ids    map[footprint]uint32
	latest map[*image.RGBA]footprint // the footprint of each poster that is currently uploaded

	holdUntil time.Time    // no real images until then (the window is being resized)
	debug     func(string) // reports what is uploaded, if set
	// cellPixels returns the size of a cell in pixels (0 if unknown): images are
	// scaled to exactly fill their cells, because terminals do not all enlarge
	// an image that is smaller than its placement (Ghostty shows it at its own
	// size, so a small poster came out smaller than the others).
	cellPixels func() (w, h int)

	// background, if set, runs prepare off the UI goroutine and then ready on
	// it. Scaling and encoding a poster takes about 200 ms; done in Draw, every
	// step to a title with a poster not uploaded yet waited that long. Without
	// it (tests) posters are prepared at once.
	background func(prepare func() prepared, ready func(prepared))
	preparing  map[footprint]bool
	done       map[footprint]prepared // prepared, not uploaded yet
}

// prepared is a poster encoded for upload.
type prepared struct {
	encoded string // base64 PNG
	w, h    int    // its size in pixels
	err     error
}

// maxPrepared bounds the prepared posters that wait for a draw (the
// selection moved on before they were ready).
const maxPrepared = 8

// settleDelay is how long the window size must stay the same before real
// images are shown again. While a window is dragged to another size, tens of
// sizes go by; uploading and placing images for each of them races with the
// terminal's own resizing, and the result can be a placement that does not
// match the cells (a small poster in a big area, or a corner of a big one).
const settleDelay = 300 * time.Millisecond

// hold stops real images being drawn until the window has been left alone for settleDelay.
func (k *kittyGraphics) hold() { k.holdUntil = time.Now().Add(settleDelay) }

// settling reports whether the window size changed so recently that blocks are drawn instead.
func (k *kittyGraphics) settling() bool { return time.Now().Before(k.holdUntil) }

// footprint identifies an uploaded image: a poster, the part of it that was
// sent, and the cells it covers.
type footprint struct {
	img        *image.RGBA
	rect       image.Rectangle
	cols, rows int
}

func newKittyGraphics(write func([]byte) error) *kittyGraphics {
	return &kittyGraphics{write: write, nextID: 1, ids: map[footprint]uint32{}, latest: map[*image.RGBA]footprint{},
		preparing: map[footprint]bool{}, done: map[footprint]prepared{}}
}

// kittyAPC builds one graphics command: ESC _ G <keys> ; <payload> ESC \.
func kittyAPC(keys string, payload []byte) []byte {
	var b bytes.Buffer
	b.WriteString("\x1b_G")
	b.WriteString(keys)
	if len(payload) > 0 {
		b.WriteByte(';')
		b.Write(payload)
	}
	b.WriteString("\x1b\\")
	return b.Bytes()
}

// show makes sure the part rect of the image is uploaded and placed over
// cols×rows cells, and returns the id the placeholder cells must carry, or 0
// while the poster is still being prepared in the background (show is called
// again on the redraw when it is ready). The same footprint is reused; a
// different one for the same poster replaces the old upload (which is removed
// from the terminal).
func (k *kittyGraphics) show(img *image.RGBA, rect image.Rectangle, cols, rows int) (uint32, error) {
	key := footprint{img, rect, cols, rows}
	if id, ok := k.ids[key]; ok {
		return id, nil
	}
	cw, ch := 0, 0
	if k.cellPixels != nil {
		cw, ch = k.cellPixels()
	}
	prepare := func() prepared { return encodePoster(img, rect, cols*cw, rows*ch) }
	p, ok := k.done[key]
	switch {
	case ok:
		delete(k.done, key)
	case k.background == nil:
		p = prepare()
	default:
		if !k.preparing[key] {
			k.preparing[key] = true
			k.background(prepare, func(p prepared) {
				delete(k.preparing, key)
				if len(k.done) >= maxPrepared {
					clear(k.done)
				}
				k.done[key] = p
			})
		}
		return 0, nil
	}
	if p.err != nil {
		return 0, p.err
	}
	return k.upload(key, p)
}

// encodePoster scales the part rect of img to w×h pixels (if both are known)
// and encodes it. Safe for use from any goroutine.
func encodePoster(img *image.RGBA, rect image.Rectangle, w, h int) prepared {
	var src image.Image = img.SubImage(rect)
	if w > 0 && h > 0 {
		src = scaleImage(src, w, h)
	}
	var data bytes.Buffer
	enc := png.Encoder{CompressionLevel: png.BestSpeed}
	if err := enc.Encode(&data, src); err != nil {
		return prepared{err: err}
	}
	return prepared{encoded: base64.StdEncoding.EncodeToString(data.Bytes()), w: src.Bounds().Dx(), h: src.Bounds().Dy()}
}

// upload sends a prepared poster and places it.
func (k *kittyGraphics) upload(key footprint, p prepared) (uint32, error) {
	img, cols, rows := key.img, key.cols, key.rows
	encoded := p.encoded
	id := k.nextID
	for first, rest := true, encoded; len(rest) > 0; first = false {
		n := min(len(rest), kittyChunk)
		more := 0
		if n < len(rest) {
			more = 1
		}
		keys := fmt.Sprintf("m=%d", more)
		if first {
			// Transmit only (a=t), PNG data (f=100) directly in the command (t=d),
			// no replies (q=2): a reply would arrive as keyboard input.
			keys = fmt.Sprintf("a=t,f=100,t=d,i=%d,q=2,m=%d", id, more)
		}
		if err := k.write(kittyAPC(keys, []byte(rest[:n]))); err != nil {
			return 0, err
		}
		rest = rest[n:]
	}
	// U=1: a virtual placement, shown wherever placeholder cells refer to the id.
	if err := k.write(kittyAPC(fmt.Sprintf("a=p,U=1,i=%d,c=%d,r=%d,q=2", id, cols, rows), nil)); err != nil {
		return 0, err
	}
	k.nextID++
	if old, ok := k.latest[img]; ok {
		if oldID, ok := k.ids[old]; ok {
			_ = k.write(kittyAPC(fmt.Sprintf("a=d,d=I,i=%d,q=2", oldID), nil))
			delete(k.ids, old)
		}
	}
	k.ids[key] = id
	k.latest[img] = key
	if k.debug != nil {
		k.debug(fmt.Sprintf("graphics: id %d, poster %dx%d px cropped to %v, sent as %dx%d px, placed over %dx%d cells",
			id, img.Bounds().Dx(), img.Bounds().Dy(), key.rect, p.w, p.h, cols, rows))
	}
	return id, nil
}

// deleteAll removes every image uploaded so far from the terminal. The next
// draw uploads what it needs again, which is also how a resize is handled.
func (k *kittyGraphics) deleteAll() {
	for _, id := range k.ids {
		_ = k.write(kittyAPC(fmt.Sprintf("a=d,d=I,i=%d,q=2", id), nil))
	}
	k.ids = map[footprint]uint32{}
	k.latest = map[*image.RGBA]footprint{}
	clear(k.done) // posters in preparation still arrive, but are not used without a draw asking for them
}

// placeholderCell returns what to draw in the cell at (row, col) of the image
// with the given id: the placeholder character, its combining characters and
// the style that carries the id.
func placeholderCell(id uint32, row, col int) (main rune, comb []rune, style tcell.Style) {
	comb = []rune{kittyDiacritics[row], kittyDiacritics[col]}
	fg := tcell.NewRGBColor(int32(id>>16&0xFF), int32(id>>8&0xFF), int32(id&0xFF))
	// The background shows wherever the picture does not cover the cell (a
	// sliver at the edge): use the panes' own color so it is not visible.
	return placeholderRune, comb, tcell.StyleDefault.Foreground(fg).Background(tview.Styles.PrimitiveBackgroundColor)
}

// maxPlaceholders is how many rows and columns can be numbered.
const maxPlaceholders = len(kittyDiacritics)

// kittyTerminal reports whether the environment looks like a terminal that
// supports Kitty graphics with Unicode placeholders: kitty and Ghostty, not
// inside tmux or screen (which swallow the graphics commands).
func kittyTerminal(env func(string) string) bool {
	term := strings.ToLower(env("TERM"))
	if env("TMUX") != "" || env("STY") != "" || strings.HasPrefix(term, "screen") || strings.HasPrefix(term, "tmux") {
		return false
	}
	switch {
	case strings.Contains(term, "kitty"), strings.Contains(term, "ghostty"):
		return true
	case strings.EqualFold(env("TERM_PROGRAM"), "ghostty"), env("GHOSTTY_RESOURCES_DIR") != "", env("KITTY_WINDOW_ID") != "":
		return true
	}
	return false
}

// scaleImage resizes src to w×h pixels, smoothly (bilinear) both up and down.
// Sizes beyond 4000 pixels are not made (a misreported cell size).
func scaleImage(src image.Image, w, h int) image.Image {
	b := src.Bounds()
	if w <= 0 || h <= 0 || w > 4000 || h > 4000 || (w == b.Dx() && h == b.Dy()) {
		return src
	}
	dst := image.NewRGBA(image.Rect(0, 0, w, h))
	xdraw.BiLinear.Scale(dst, dst.Bounds(), src, b, xdraw.Src, nil)
	return dst
}