package main

import (
	"fmt"
	"sort"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"
)

// ---------------------------------------------------------------------------
// Colour and attribute model
// ---------------------------------------------------------------------------

// colorMode says how a termColor should be resolved against the theme palette.
type colorMode uint8

const (
	colDefault colorMode = iota // the theme's foreground / background
	colIndexed                  // 0-255: the 16 ANSI colours, the 6x6x6 cube, the greys
	colRGB                      // 24-bit truecolour, exactly as given
)

// termColor is one resolved SGR colour. It is a value type so cells compare
// with == and grids compare with reflect.DeepEqual.
type termColor struct {
	Mode colorMode `json:"mode"`
	Idx  uint8     `json:"idx"`
	R    uint8     `json:"r"`
	G    uint8     `json:"g"`
	B    uint8     `json:"b"`
}

func indexedColor(i int) termColor { return termColor{Mode: colIndexed, Idx: uint8(i)} }

func rgbColor(r, g, b int) termColor {
	return termColor{Mode: colRGB, R: uint8(r), G: uint8(g), B: uint8(b)}
}

// Character attributes. Blink, italic, strike-through and conceal are parsed
// but not rendered; see the skipped-sequence report.
const (
	attrBold uint8 = 1 << iota
	attrDim
	attrReverse
	attrUnderline
)

// cell is one character position on the grid.
type cell struct {
	Ch    rune
	FG    termColor
	BG    termColor
	Attrs uint8
}

func blankCell() cell { return cell{Ch: ' '} }

// blankWith returns an erase-fill cell carrying the current background, which
// is what a real terminal writes when ED/EL runs with a background colour set.
func blankWith(pen cell) cell { return cell{Ch: ' ', BG: pen.BG} }

// ---------------------------------------------------------------------------
// Terminal
// ---------------------------------------------------------------------------

// parseState is the ANSI byte-stream state machine's current position.
type parseState uint8

const (
	stGround parseState = iota
	stEsc
	stCSI
	stCSIIgnore // a CSI too long to buffer: swallow it up to its final byte
	stOSC       // ESC ] ... BEL | ESC \
	stOSCEsc    // saw ESC inside an OSC string
	stCharset
)

// terminal is an in-memory character grid plus the ANSI parser that drives it.
// There is no scrollback: lines that scroll off the top are gone.
type terminal struct {
	rows, cols int
	cells      []cell // len == rows*cols, row-major
	cx, cy     int
	pen        cell // current SGR state; Ch is unused
	wrapNext   bool // cursor is parked in the last column with a wrap pending

	state   parseState
	csi     []byte
	utf8buf []byte

	printed   int64
	skipped   map[string]int
	skipOrder []string
}

const (
	maxGridCols = 1000
	maxGridRows = 1000
	tabWidth    = 8
)

func newTerminal(cols, rows int) (*terminal, error) {
	if cols <= 0 || rows <= 0 {
		return nil, fmt.Errorf("grid must be at least 1x1, got %dx%d", cols, rows)
	}
	if cols > maxGridCols || rows > maxGridRows {
		return nil, fmt.Errorf("grid %dx%d is larger than the %dx%d limit", cols, rows, maxGridCols, maxGridRows)
	}
	t := &terminal{
		rows:    rows,
		cols:    cols,
		cells:   make([]cell, rows*cols),
		pen:     blankCell(),
		skipped: map[string]int{},
	}
	t.fillAll(blankCell())
	return t, nil
}

func (t *terminal) fillAll(c cell) {
	for i := range t.cells {
		t.cells[i] = c
	}
}

func (t *terminal) at(y, x int) *cell { return &t.cells[y*t.cols+x] }

func (t *terminal) note(what string) {
	if _, ok := t.skipped[what]; !ok {
		t.skipOrder = append(t.skipOrder, what)
	}
	t.skipped[what]++
}

// skippedReport lists every unsupported sequence encountered, most frequent
// first, in a stable order.
func (t *terminal) skippedReport() []string {
	out := append([]string(nil), t.skipOrder...)
	sort.SliceStable(out, func(i, j int) bool { return t.skipped[out[i]] > t.skipped[out[j]] })
	lines := make([]string, 0, len(out))
	for _, k := range out {
		lines = append(lines, fmt.Sprintf("%-28s x%d", k, t.skipped[k]))
	}
	return lines
}

func (t *terminal) skippedTotal() int {
	n := 0
	for _, v := range t.skipped {
		n += v
	}
	return n
}

// clone produces an independent deep copy, so a replay can be checkpointed and
// resumed without the resumed run observing the original.
func (t *terminal) clone() *terminal {
	c := &terminal{
		rows: t.rows, cols: t.cols,
		cells:    append([]cell(nil), t.cells...),
		cx:       t.cx,
		cy:       t.cy,
		pen:      t.pen,
		wrapNext: t.wrapNext,
		state:    t.state,
		csi:      append([]byte(nil), t.csi...),
		utf8buf:  append([]byte(nil), t.utf8buf...),
		printed:  t.printed,
		skipped:  map[string]int{},
	}
	for k, v := range t.skipped {
		c.skipped[k] = v
	}
	c.skipOrder = append([]string(nil), t.skipOrder...)
	return c
}

// gridSnapshot is the comparable state of the visible screen.
type gridSnapshot struct {
	Rows, Cols int
	CX, CY     int
	Cells      []cell
}

func (t *terminal) snapshot() gridSnapshot {
	return gridSnapshot{
		Rows: t.rows, Cols: t.cols,
		CX: t.cx, CY: t.cy,
		Cells: append([]cell(nil), t.cells...),
	}
}

// lineText renders one row as plain text with trailing blanks removed.
func (t *terminal) lineText(y int) string {
	if y < 0 || y >= t.rows {
		return ""
	}
	var b strings.Builder
	for x := 0; x < t.cols; x++ {
		b.WriteRune(t.at(y, x).Ch)
	}
	return strings.TrimRight(b.String(), " ")
}

// text renders the whole grid, one line per row, trailing blanks removed.
func (t *terminal) text() string {
	lines := make([]string, t.rows)
	for y := 0; y < t.rows; y++ {
		lines[y] = t.lineText(y)
	}
	return strings.Join(lines, "\n")
}

// ---------------------------------------------------------------------------
// Cursor and screen primitives
// ---------------------------------------------------------------------------

func clampInt(v, lo, hi int) int {
	if v < lo {
		return lo
	}
	if v > hi {
		return hi
	}
	return v
}

func (t *terminal) scrollUp(n int) {
	if n <= 0 {
		return
	}
	if n >= t.rows {
		t.fillAll(blankWith(t.pen))
		return
	}
	copy(t.cells, t.cells[n*t.cols:])
	fill := blankWith(t.pen)
	for i := (t.rows - n) * t.cols; i < len(t.cells); i++ {
		t.cells[i] = fill
	}
}

// lineFeed moves down one row, scrolling the grid when it runs off the bottom.
// The column is left alone; newline is the one that also returns the carriage.
func (t *terminal) lineFeed() {
	t.wrapNext = false
	t.cy++
	if t.cy >= t.rows {
		t.scrollUp(t.cy - t.rows + 1)
		t.cy = t.rows - 1
	}
}

// newline is LF with the carriage return the tty driver would have added.
//
// This is the one place CaptureStudio deliberately differs from a bare VT100.
// SessionForge captures PIPES, not a PTY, so its casts contain the bytes the
// program wrote: a lone "\n". On a real terminal that program's "\n" would have
// passed through the kernel line discipline with ONLCR set, which turns it into
// CR+LF before the terminal ever sees it. Replaying pipe-captured output with
// strict LF-means-down-only semantics produces the famous descending staircase
// instead of a screen. So LF, VT and FF all return the carriage here, exactly
// as the tty driver would have done for a program printing to a terminal.
func (t *terminal) newline() {
	t.carriageReturn()
	t.lineFeed()
}

func (t *terminal) carriageReturn() {
	t.cx = 0
	t.wrapNext = false
}

func (t *terminal) backspace() {
	t.wrapNext = false
	if t.cx > 0 {
		t.cx--
	}
}

// tab advances to the next multiple-of-8 column. Tab stops are fixed; HTS and
// TBC are not implemented.
func (t *terminal) tab() {
	t.wrapNext = false
	next := (t.cx/tabWidth + 1) * tabWidth
	if next > t.cols-1 {
		next = t.cols - 1
	}
	t.cx = next
}

// putRune writes one character at the cursor, honouring deferred wrapping: a
// character written into the last column parks the cursor there and only wraps
// when the NEXT character arrives. That is what xterm does, and it is why
// printing exactly `cols` characters does not leave a blank line behind.
func (t *terminal) putRune(r rune) {
	if t.wrapNext {
		t.cx = 0
		t.lineFeed()
	}
	c := t.pen
	c.Ch = r
	*t.at(t.cy, t.cx) = c
	t.printed++
	if t.cx+1 >= t.cols {
		t.wrapNext = true
	} else {
		t.cx++
	}
}

// eraseRange clears cells [from, to) in row-major order.
func (t *terminal) eraseRange(from, to int) {
	fill := blankWith(t.pen)
	from = clampInt(from, 0, len(t.cells))
	to = clampInt(to, 0, len(t.cells))
	for i := from; i < to; i++ {
		t.cells[i] = fill
	}
}

func (t *terminal) eraseDisplay(mode int) {
	pos := t.cy*t.cols + t.cx
	switch mode {
	case 0:
		t.eraseRange(pos, len(t.cells))
	case 1:
		t.eraseRange(0, pos+1)
	case 2, 3:
		// Mode 3 also clears scrollback on a real terminal. There is none here,
		// so it behaves exactly like mode 2.
		t.eraseRange(0, len(t.cells))
	default:
		t.note(fmt.Sprintf("CSI %dJ (erase mode)", mode))
	}
	t.wrapNext = false
}

func (t *terminal) eraseLine(mode int) {
	start := t.cy * t.cols
	switch mode {
	case 0:
		t.eraseRange(start+t.cx, start+t.cols)
	case 1:
		t.eraseRange(start, start+t.cx+1)
	case 2:
		t.eraseRange(start, start+t.cols)
	default:
		t.note(fmt.Sprintf("CSI %dK (erase mode)", mode))
	}
	t.wrapNext = false
}

// ---------------------------------------------------------------------------
// Byte stream parser
// ---------------------------------------------------------------------------

// Write feeds bytes to the emulator. It never returns an error: unsupported
// input is counted and skipped, never allowed to corrupt the grid.
func (t *terminal) Write(p []byte) (int, error) {
	for _, b := range p {
		t.feed(b)
	}
	return len(p), nil
}

func (t *terminal) feed(b byte) {
	switch t.state {
	case stGround:
		t.feedGround(b)
	case stEsc:
		t.feedEsc(b)
	case stCSI:
		t.feedCSI(b)
	case stCSIIgnore:
		// Still inside an abandoned CSI. Swallow everything up to and including
		// its final byte, so its parameters never leak onto the screen as text.
		if b >= 0x40 && b <= 0x7e {
			t.state = stGround
		}
	case stOSC:
		switch b {
		case 0x07: // BEL terminates
			t.state = stGround
		case 0x1b:
			t.state = stOSCEsc
		}
	case stOSCEsc:
		// ESC \ (String Terminator) ends the OSC; anything else resumes it.
		t.state = stOSC
		if b == '\\' {
			t.state = stGround
		}
	case stCharset:
		// ESC ( B and friends: designate a character set. Only one byte long and
		// we always behave as if US-ASCII is selected.
		t.note("ESC ( charset select")
		t.state = stGround
	}
}

func (t *terminal) feedGround(b byte) {
	if b >= 0x80 {
		t.feedUTF8(b)
		return
	}
	// A stray high byte sequence is abandoned as soon as ASCII resumes.
	if len(t.utf8buf) > 0 {
		t.flushBadUTF8()
	}
	switch {
	case b == 0x1b:
		t.state = stEsc
		t.csi = t.csi[:0]
	case b == 0x07: // BEL: audible only, nothing to draw
	case b == 0x08:
		t.backspace()
	case b == 0x09:
		t.tab()
	case b == 0x0a, b == 0x0b, b == 0x0c:
		// LF, VT and FF: down one line AND back to column 0. See newline().
		t.newline()
	case b == 0x0d:
		t.carriageReturn()
	case b < 0x20 || b == 0x7f:
		t.note(fmt.Sprintf("C0 control 0x%02x", b))
	default:
		t.putRune(rune(b))
	}
}

// feedUTF8 accumulates a multi-byte rune. The bitmap font only covers ASCII, so
// a decoded non-ASCII rune still occupies exactly one cell and is drawn with
// the missing-glyph box; it is stored faithfully so text extraction is correct.
func (t *terminal) feedUTF8(b byte) {
	t.utf8buf = append(t.utf8buf, b)
	if !utf8.FullRune(t.utf8buf) {
		if len(t.utf8buf) >= utf8.UTFMax {
			t.flushBadUTF8()
		}
		return
	}
	r, size := utf8.DecodeRune(t.utf8buf)
	if r == utf8.RuneError && size <= 1 {
		t.flushBadUTF8()
		return
	}
	t.utf8buf = t.utf8buf[:0]
	if unicode.IsControl(r) {
		t.note("C1/unicode control")
		return
	}
	t.putRune(r)
}

func (t *terminal) flushBadUTF8() {
	t.utf8buf = t.utf8buf[:0]
	t.note("invalid UTF-8 byte")
	t.putRune(unicode.ReplacementChar)
}

func (t *terminal) feedEsc(b byte) {
	switch b {
	case '[':
		t.state = stCSI
		t.csi = t.csi[:0]
	case ']':
		t.state = stOSC
	case '(', ')', '*', '+':
		t.state = stCharset
	case 'P', 'X', '^', '_':
		// DCS / SOS / PM / APC all carry a string terminated the same way as OSC.
		t.note("ESC " + string(rune(b)) + " string")
		t.state = stOSC
	case 'c':
		t.reset()
		t.state = stGround
	case 'M': // Reverse index: not implemented, would need a scroll region.
		t.note("ESC M reverse index")
		t.state = stGround
	case '\\':
		t.state = stGround
	default:
		t.note("ESC " + string(rune(b)))
		t.state = stGround
	}
}

func (t *terminal) reset() {
	t.pen = blankCell()
	t.fillAll(blankCell())
	t.cx, t.cy = 0, 0
	t.wrapNext = false
}

const maxCSILen = 64

func (t *terminal) feedCSI(b byte) {
	if b >= 0x40 && b <= 0x7e { // final byte
		t.execCSI(string(t.csi), b)
		t.csi = t.csi[:0]
		t.state = stGround
		return
	}
	if len(t.csi) >= maxCSILen {
		// Runaway sequence: stop buffering it, but keep swallowing bytes until
		// its final byte arrives so the parameters do not spill onto the grid.
		t.note("over-long CSI sequence")
		t.csi = t.csi[:0]
		t.state = stCSIIgnore
		return
	}
	t.csi = append(t.csi, b)
}

// csiParams splits a CSI parameter string on ';'. An empty field means "use the
// default", which callers express by passing the default to paramAt.
func csiParams(s string) []string {
	if s == "" {
		return nil
	}
	return strings.Split(s, ";")
}

func paramAt(ps []string, i, def int) int {
	if i >= len(ps) {
		return def
	}
	f := ps[i]
	if c := strings.IndexByte(f, ':'); c >= 0 {
		f = f[:c]
	}
	if f == "" {
		return def
	}
	n, err := strconv.Atoi(f)
	if err != nil || n < 0 {
		return def
	}
	return n
}

func (t *terminal) execCSI(params string, final byte) {
	// Private / experimental sequences (?, >, <, =) are all out of scope: they
	// set modes we do not model (cursor visibility, bracketed paste, mouse).
	if params != "" && (params[0] == '?' || params[0] == '>' || params[0] == '<' || params[0] == '=') {
		t.note(fmt.Sprintf("CSI %c...%c (private mode)", params[0], final))
		return
	}
	ps := csiParams(params)
	switch final {
	case 'A': // CUU
		t.cy = clampInt(t.cy-maxInt(1, paramAt(ps, 0, 1)), 0, t.rows-1)
		t.wrapNext = false
	case 'B': // CUD
		t.cy = clampInt(t.cy+maxInt(1, paramAt(ps, 0, 1)), 0, t.rows-1)
		t.wrapNext = false
	case 'C': // CUF
		t.cx = clampInt(t.cx+maxInt(1, paramAt(ps, 0, 1)), 0, t.cols-1)
		t.wrapNext = false
	case 'D': // CUB
		t.cx = clampInt(t.cx-maxInt(1, paramAt(ps, 0, 1)), 0, t.cols-1)
		t.wrapNext = false
	case 'E': // CNL
		t.cy = clampInt(t.cy+maxInt(1, paramAt(ps, 0, 1)), 0, t.rows-1)
		t.cx = 0
		t.wrapNext = false
	case 'F': // CPL
		t.cy = clampInt(t.cy-maxInt(1, paramAt(ps, 0, 1)), 0, t.rows-1)
		t.cx = 0
		t.wrapNext = false
	case 'G', '`': // CHA / HPA, 1-based column
		t.cx = clampInt(paramAt(ps, 0, 1)-1, 0, t.cols-1)
		t.wrapNext = false
	case 'd': // VPA, 1-based row
		t.cy = clampInt(paramAt(ps, 0, 1)-1, 0, t.rows-1)
		t.wrapNext = false
	case 'H', 'f': // CUP / HVP, both parameters 1-based
		t.cy = clampInt(paramAt(ps, 0, 1)-1, 0, t.rows-1)
		t.cx = clampInt(paramAt(ps, 1, 1)-1, 0, t.cols-1)
		t.wrapNext = false
	case 'J':
		t.eraseDisplay(paramAt(ps, 0, 0))
	case 'K':
		t.eraseLine(paramAt(ps, 0, 0))
	case 'm':
		t.execSGR(ps)
	default:
		t.note(fmt.Sprintf("CSI %c", final))
	}
}

func maxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// execSGR applies Select Graphic Rendition. Both the classic semicolon form
// (38;5;n and 38;2;r;g;b) and the colon sub-parameter form (38:5:n,
// 38:2::r:g:b) are accepted, because real programs emit both.
func (t *terminal) execSGR(ps []string) {
	if len(ps) == 0 {
		t.pen = blankCell()
		return
	}
	for i := 0; i < len(ps); i++ {
		if strings.ContainsRune(ps[i], ':') {
			t.sgrColonForm(ps[i])
			continue
		}
		n := paramAt(ps, i, 0)
		switch {
		case n == 0:
			t.pen = blankCell()
		case n == 1:
			t.pen.Attrs |= attrBold
		case n == 2:
			t.pen.Attrs |= attrDim
		case n == 4:
			t.pen.Attrs |= attrUnderline
		case n == 7:
			t.pen.Attrs |= attrReverse
		case n == 21, n == 22:
			t.pen.Attrs &^= attrBold | attrDim
		case n == 24:
			t.pen.Attrs &^= attrUnderline
		case n == 27:
			t.pen.Attrs &^= attrReverse
		case n >= 30 && n <= 37:
			t.pen.FG = indexedColor(n - 30)
		case n == 39:
			t.pen.FG = termColor{}
		case n >= 40 && n <= 47:
			t.pen.BG = indexedColor(n - 40)
		case n == 49:
			t.pen.BG = termColor{}
		case n >= 90 && n <= 97:
			t.pen.FG = indexedColor(n - 90 + 8)
		case n >= 100 && n <= 107:
			t.pen.BG = indexedColor(n - 100 + 8)
		case n == 38 || n == 48:
			used, col, ok := parseExtendedColor(ps, i)
			if !ok {
				t.note(fmt.Sprintf("SGR %d (extended colour)", n))
				return // the rest of the parameter list is no longer trustworthy
			}
			if n == 38 {
				t.pen.FG = col
			} else {
				t.pen.BG = col
			}
			i += used
		default:
			t.note(fmt.Sprintf("SGR %d", n))
		}
	}
}

// parseExtendedColor reads ps[i] == "38"/"48" plus its following selector and
// operands, returning how many EXTRA fields it consumed.
func parseExtendedColor(ps []string, i int) (int, termColor, bool) {
	if i+1 >= len(ps) {
		return 0, termColor{}, false
	}
	switch paramAt(ps, i+1, -1) {
	case 5:
		if i+2 >= len(ps) {
			return 0, termColor{}, false
		}
		v := paramAt(ps, i+2, -1)
		if v < 0 || v > 255 {
			return 0, termColor{}, false
		}
		return 2, indexedColor(v), true
	case 2:
		if i+4 >= len(ps) {
			return 0, termColor{}, false
		}
		r, g, b := paramAt(ps, i+2, -1), paramAt(ps, i+3, -1), paramAt(ps, i+4, -1)
		if r < 0 || r > 255 || g < 0 || g > 255 || b < 0 || b > 255 {
			return 0, termColor{}, false
		}
		return 4, rgbColor(r, g, b), true
	}
	return 0, termColor{}, false
}

// sgrColonForm handles a single parameter that carries its own sub-parameters,
// e.g. "38:5:196" or "38:2::12:34:56" (the empty field is a colour space id).
func (t *terminal) sgrColonForm(field string) {
	parts := strings.Split(field, ":")
	head := parts[0]
	n, err := strconv.Atoi(head)
	if err != nil {
		t.note("SGR malformed sub-parameter")
		return
	}
	if n != 38 && n != 48 {
		// Underline styles (4:3) and similar: the base attribute still applies.
		t.execSGR([]string{head})
		return
	}
	nums := make([]int, 0, len(parts))
	for _, p := range parts[1:] {
		if p == "" {
			continue // colour-space identifier, per ITU T.416
		}
		v, err := strconv.Atoi(p)
		if err != nil {
			t.note("SGR malformed sub-parameter")
			return
		}
		nums = append(nums, v)
	}
	var col termColor
	switch {
	case len(nums) == 2 && nums[0] == 5 && nums[1] >= 0 && nums[1] <= 255:
		col = indexedColor(nums[1])
	case len(nums) == 4 && nums[0] == 2 && inByte(nums[1]) && inByte(nums[2]) && inByte(nums[3]):
		col = rgbColor(nums[1], nums[2], nums[3])
	default:
		t.note(fmt.Sprintf("SGR %d (extended colour)", n))
		return
	}
	if n == 38 {
		t.pen.FG = col
	} else {
		t.pen.BG = col
	}
}

func inByte(v int) bool { return v >= 0 && v <= 255 }