package main

import (
	"bytes"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/racingmars/go3270"
)

func indexTelnetEOR(b []byte) int {
	for i := 0; i+1 < len(b); i++ {
		if b[i] == IAC {
			switch b[i+1] {
			case IAC:
				i++ // skip second 0xFF of escaped IAC IAC pair
			case EOR:
				return i
			}
		}
	}
	return -1
}

func lastIndexTelnetEOR(b []byte) int {
	last := -1
	for i := 0; i+1 < len(b); i++ {
		if b[i] == IAC {
			switch b[i+1] {
			case IAC:
				i++
			case EOR:
				last = i
				i++
			}
		}
	}
	return last
}

// negotiate RFC 2355 TN3270E or RFC 1576 TN3270 options with an upstream host
func negotiateUpstreamTelnet(conn net.Conn, termType string) ([]byte, bool, bool, error) {
	_ = conn.SetDeadline(time.Now().Add(7 * time.Second))
	defer conn.SetDeadline(time.Time{})

	var isTN3270E bool
	var devTypeSent bool
	var functionsDone bool
	var agreedBindImage bool

	var ourBinary, hisBinary bool
	var ourEOR, hisEOR bool
	var ttypeSent bool

	buf := make([]byte, 8192)
	var initial3270Data []byte

	for {
		isTN3270EDone := isTN3270E && devTypeSent && functionsDone
		isTN3270Done := !isTN3270E && ourBinary && hisBinary && ourEOR && hisEOR && ttypeSent

		if isTN3270EDone || isTN3270Done || len(initial3270Data) > 0 {
			break
		}

		n, err := conn.Read(buf)
		if err != nil {
			if (isTN3270E && functionsDone) || (ourBinary && hisBinary && ourEOR && hisEOR) {
				break
			}
			return nil, false, false, fmt.Errorf("read during upstream negotiation: %w", err)
		}
		data := buf[:n]

		var replies []byte
		i := 0
		for i < len(data) {
			if data[i] == IAC && i+1 < len(data) {
				cmd := data[i+1]
				if cmd == SB {
					seIdx := bytes.Index(data[i:], []byte{IAC, SE})
					if seIdx != -1 {
						sbPayload := data[i+2 : i+seIdx]
						if len(sbPayload) >= 2 && sbPayload[0] == OPT_TERMINAL_TYPE && sbPayload[1] == TTYPE_SEND {
							replies = append(replies, IAC, SB, OPT_TERMINAL_TYPE, TTYPE_IS)
							replies = append(replies, []byte(termType)...)
							replies = append(replies, IAC, SE)
							ttypeSent = true
						} else if len(sbPayload) >= 3 && sbPayload[0] == OPT_TN3270E {
							switch sbPayload[1] {
							case TN3270E_SEND:
								if sbPayload[2] == TN3270E_DEVICE_TYPE {
									// RFC 2355: TN3270E device-type must NOT contain -E suffix
									devType := strings.TrimSuffix(termType, "-E")
									replies = append(replies, IAC, SB, OPT_TN3270E, TN3270E_DEVICE_TYPE, TN3270E_REQUEST)
									replies = append(replies, []byte(devType)...)
									replies = append(replies, IAC, SE)
									devTypeSent = true
								}
							case TN3270E_DEVICE_TYPE:
								switch sbPayload[2] {
								case TN3270E_IS:
									replies = append(replies, IAC, SB, OPT_TN3270E, TN3270E_FUNCTIONS, TN3270E_REQUEST, TN3270E_FUNC_RESPONSES, TN3270E_FUNC_BIND_IMAGE, IAC, SE)
								case TN3270E_REJECT:
									// Server rejected TN3270E device type -> fall back to basic TN3270
									isTN3270E = false
									replies = append(replies, IAC, WONT, OPT_TN3270E, IAC, DONT, OPT_TN3270E)
								}
							case TN3270E_FUNCTIONS:
								if sbPayload[2] == TN3270E_REQUEST || sbPayload[2] == TN3270E_IS {
									funcs := sbPayload[3:]
									if bytes.Contains(funcs, []byte{TN3270E_FUNC_BIND_IMAGE}) {
										agreedBindImage = true
									}
									if sbPayload[2] == TN3270E_REQUEST {
										replies = append(replies, IAC, SB, OPT_TN3270E, TN3270E_FUNCTIONS, TN3270E_IS)
										replies = append(replies, funcs...)
										replies = append(replies, IAC, SE)
									}
									functionsDone = true
								}
							}
						}
						i += seIdx + 2
						continue
					}
				}
				if i+2 < len(data) {
					opt := data[i+2]
					switch cmd {
					case DO:
						switch opt {
						case OPT_TN3270E:
							replies = append(replies, IAC, WILL, opt)
							isTN3270E = true
						case OPT_BINARY:
							replies = append(replies, IAC, WILL, opt)
							ourBinary = true
						case OPT_EOR:
							replies = append(replies, IAC, WILL, opt)
							ourEOR = true
						case OPT_TIME_MARK, OPT_SUPPRESS_GA, OPT_TERMINAL_TYPE:
							replies = append(replies, IAC, WILL, opt)
						default:
							replies = append(replies, IAC, WONT, opt)
						}
						i += 3
						continue
					case WILL:
						switch opt {
						case OPT_TN3270E:
							replies = append(replies, IAC, DO, opt)
							isTN3270E = true
						case OPT_BINARY:
							replies = append(replies, IAC, DO, opt)
							hisBinary = true
						case OPT_EOR:
							replies = append(replies, IAC, DO, opt)
							hisEOR = true
						case OPT_SUPPRESS_GA, OPT_TERMINAL_TYPE:
							replies = append(replies, IAC, DO, opt)
						default:
							replies = append(replies, IAC, DONT, opt)
						}
						i += 3
						continue
					case DONT:
						switch opt {
						case OPT_TN3270E:
							isTN3270E = false
						case OPT_BINARY:
							ourBinary = false
						case OPT_EOR:
							ourEOR = false
						}
						replies = append(replies, IAC, WONT, opt)
						i += 3
						continue
					case WONT:
						switch opt {
						case OPT_TN3270E:
							isTN3270E = false
						case OPT_BINARY:
							hisBinary = false
						case OPT_EOR:
							hisEOR = false
						}
						replies = append(replies, IAC, DONT, opt)
						i += 3
						continue
					}
				}
				// 2-byte Telnet commands like NOP (0xF1), GA (0xF9)
				if cmd >= NOP && cmd <= GA {
					i += 2
					continue
				}
			}

			// Non-negotiation bytes (initial 3270 data stream)
			initial3270Data = append(initial3270Data, data[i:]...)
			break
		}

		if len(replies) > 0 {
			if _, err := conn.Write(replies); err != nil {
				return nil, false, false, err
			}
		}
	}

	finalTN3270E := isTN3270E && functionsDone

	var isBound bool
	if finalTN3270E && !agreedBindImage {
		isBound = true
	}

	// If TN3270E, parse initial3270Data records
	if finalTN3270E && len(initial3270Data) >= 5 {
		var displayable []byte
		recBuf := initial3270Data
		for len(recBuf) >= 5 {
			eorIdx := indexTelnetEOR(recBuf)
			var record []byte
			if eorIdx != -1 {
				record = recBuf[:eorIdx+2]
				recBuf = recBuf[eorIdx+2:]
			} else {
				record = recBuf
				recBuf = nil
			}
			if len(record) >= 5 {
				dataType := record[0]
				switch dataType {
				case TN3270E_DATA_BIND:
					isBound = true
				case TN3270E_DATA_UNBIND:
					isBound = false
				case TN3270E_DATA_3270, TN3270E_DATA_SSCP_LU:
					isBound = (dataType == TN3270E_DATA_3270)
					displayable = append(displayable, record[5:]...)
					if eorIdx == -1 && len(record) > 5 {
						displayable = append(displayable, IAC, EOR)
					}
				}
			}
		}
		initial3270Data = displayable
	}

	return initial3270Data, finalTN3270E, isBound, nil
}

func extractAID(packet []byte) go3270.AID {
	for i := 0; i < len(packet); {
		if packet[i] == IAC {
			if i+1 < len(packet) {
				switch packet[i+1] {
				case DO, DONT, WILL, WONT:
					i += 3
					continue
				case SB:
					seIdx := bytes.Index(packet[i:], []byte{IAC, SE})
					if seIdx != -1 {
						i += seIdx + 2
						continue
					}
				}
			}
			i += 2
			continue
		}
		return go3270.AID(packet[i])
	}
	return go3270.AIDNone
}

func is3270Query(data []byte) bool {
	if len(data) < 3 {
		return false
	}
	idx := bytes.IndexByte(data, CMD_WSF)
	if idx == -1 {
		return false
	}
	return bytes.Contains(data[idx+1:], []byte{SF_READ_PARTITION, SF_RP_QUERY})
}

func buildQueryReply(rows, cols int, isColor bool, isTN3270E bool, isBound bool) []byte {
	var reply []byte

	// 1. AID byte
	reply = append(reply, AID_QUERY_REPLY)

	// 2. Structured Field: Usable Area (0x81 0x81)
	cells := rows * cols
	usableArea := []byte{
		0x00, 0x1A, // Length = 26
		SF_QUERY_REPLY, QR_USABLE_AREA, // SFID = Usable Area
		0x00, 0x00, // Flags (12/14-bit addressing)
		byte(cols >> 8), byte(cols & 0xFF), // Width
		byte(rows >> 8), byte(rows & 0xFF), // Height
		0x01,                   // Units = inches (0.01 in)
		0x00, 0x0A, 0x00, 0x0A, // Xr, Yr
		0x00, 0x0A, 0x00, 0x0A, // Aw, Ah
		byte(cells >> 8), byte(cells & 0xFF), // Buffer size
	}
	reply = append(reply, usableArea...)

	// 3. Structured Field: Color (0x81 0x86) if color
	if isColor {
		colorSF := []byte{
			0x00, 0x12, // Length = 18
			SF_QUERY_REPLY, QR_COLOR, // SFID = Color
			0x00, 0x00, // Flags
			0x00,                  // Default color
			0x00,                  // Neutral color
			ATTR_COLOR_BLUE, 0x00, // Blue
			ATTR_COLOR_RED, 0x00, // Red
			ATTR_COLOR_PINK, 0x00, // Pink
			ATTR_COLOR_GREEN, 0x00, // Green
			ATTR_COLOR_TURQUOISE, 0x00, // Turquoise
			ATTR_COLOR_YELLOW, 0x00, // Yellow
			ATTR_COLOR_WHITE, 0x00, // Neutral/White
		}
		reply = append(reply, colorSF...)
	}

	// 4. Structured Field: Highlight (0x81 0x87)
	highlightSF := []byte{
		0x00, 0x0D, // Length = 13
		SF_QUERY_REPLY, QR_HIGHLIGHTING, // SFID = Highlight
		0x00, 0x00, // Flags
		0x00, 0x00, // Default
		ATTR_HL_BLINK, 0x00, // Blink
		ATTR_HL_REVERSE, 0x00, // Reverse Video
		ATTR_HL_UNDERSCORE, 0x00, // Underscore
	}
	reply = append(reply, highlightSF...)

	// 5. Structured Field: Implicit Partition (0x81 0xA6)
	implicitSF := []byte{
		0x00, 0x08, // Length = 8
		SF_QUERY_REPLY, QR_IMPLICIT_PARTITION, // SFID = Implicit Partition
		0x00, 0x00, // Flags
		byte(cols >> 8), byte(cols & 0xFF), // Alternate width
		byte(rows >> 8), byte(rows & 0xFF), // Alternate height
	}
	reply = append(reply, implicitSF...)

	// 6. Structured Field: Null / End of Query List
	nullSF := []byte{
		0x00, 0x04,
		SF_QUERY_REPLY, QR_NULL,
	}
	reply = append(reply, nullSF...)

	headerLen := 0
	if isTN3270E {
		headerLen = 5
	}
	escaped := make([]byte, headerLen, len(reply)*2+headerLen+2)
	for _, b := range reply {
		if b == IAC {
			escaped = append(escaped, IAC, IAC)
		} else {
			escaped = append(escaped, b)
		}
	}
	escaped = append(escaped, IAC, EOR)

	if isTN3270E {
		dataType := TN3270E_DATA_3270
		if !isBound {
			dataType = TN3270E_DATA_SSCP_LU
		}
		escaped[0] = dataType
		escaped[1] = 0x00
		escaped[2] = 0x00
		escaped[3] = 0x00
		escaped[4] = 0x00
	}

	return escaped
}

func GetTerminalMatrix(termType string, devInfo go3270.DevInfo) (rows, cols, maxBuf int) {
	if devInfo != nil {
		rows, cols = devInfo.AltDimensions()
	}
	if rows <= 0 || cols <= 0 {
		rows, cols = 24, 80
		for _, p := range strings.Split(strings.TrimSuffix(termType, ")"), "-") {
			switch p {
			case "3":
				rows, cols = 32, 80
			case "4":
				rows, cols = 43, 80
			case "5":
				rows, cols = 27, 132
			}
		}
	}
	maxBuf = min(max(rows*cols*4, 8192), 65536)
	return rows, cols, maxBuf
}

func ResolveTerminalType(origTermType string, devInfo go3270.DevInfo) string {
	rows, cols, _ := GetTerminalMatrix(origTermType, devInfo)

	model := 2
	switch {
	case rows == 32 && cols == 80:
		model = 3
	case rows == 43 && cols == 80:
		model = 4
	case rows == 27 && cols == 132:
		model = 5
	}

	family := "3279"
	if strings.Contains(strings.ToUpper(origTermType), "3278") {
		family = "3278"
	}

	return fmt.Sprintf("IBM-%s-%d-E", family, model)
}