package cli

import (
	"bytes"
	"net"
)

const (
	iac  = 255 // Interpret As Command
	dont = 254
	do   = 253
	wont = 252
	will = 251
	sb   = 250
	se   = 240

	optEcho             = 1
	optSuppressGoAhead  = 3
	optTerminalType     = 24
	optWindowSize       = 31
	optLinemode         = 34
)

// TelnetConn wraps a net.Conn and handles RFC 854 Telnet negotiations
// while presenting a clean stream to terminal.NewTerminal.
type TelnetConn struct {
	net.Conn
	subnegotiation bool
}

func NewTelnetConn(conn net.Conn) *TelnetConn {
	tc := &TelnetConn{Conn: conn}
	tc.initNegotiation()
	return tc
}

func (tc *TelnetConn) initNegotiation() {
	// Request client to suppress go-ahead and handle character-at-a-time / local echo off
	// IAC WILL ECHO, IAC WILL SUPPRESS_GO_AHEAD, IAC DONT LINEMODE
	initSeq := []byte{
		iac, will, optEcho,
		iac, will, optSuppressGoAhead,
		iac, dont, optLinemode,
	}
	_, _ = tc.Conn.Write(initSeq)
}

func (tc *TelnetConn) Read(p []byte) (n int, err error) {
	buf := make([]byte, len(p))
	rawN, err := tc.Conn.Read(buf)
	if err != nil {
		return 0, err
	}

	var clean bytes.Buffer
	i := 0
	for i < rawN {
		b := buf[i]
		if b == iac {
			if i+1 >= rawN {
				break
			}
			cmd := buf[i+1]
			switch cmd {
			case iac: // Escaped 255
				clean.WriteByte(iac)
				i += 2
			case will, wont, do, dont:
				if i+2 < rawN {
					opt := buf[i+2]
					tc.handleOption(cmd, opt)
					i += 3
				} else {
					i = rawN
				}
			case sb: // Subnegotiation start
				tc.subnegotiation = true
				i += 2
			case se: // Subnegotiation end
				tc.subnegotiation = false
				i += 2
			default:
				i += 2
			}
			continue
		}

		if tc.subnegotiation {
			i++
			continue
		}

		// Normalize CR NUL or CR LF into LF for standard terminal handling
		if b == '\r' {
			if i+1 < rawN && (buf[i+1] == '\n' || buf[i+1] == 0) {
				clean.WriteByte('\n')
				i += 2
				continue
			}
			clean.WriteByte('\n')
			i++
			continue
		}

		clean.WriteByte(b)
		i++
	}

	res := clean.Bytes()
	copy(p, res)
	return len(res), nil
}

func (tc *TelnetConn) handleOption(cmd, opt byte) {
	// Respond to DO ECHO or DO SUPPRESS_GO_AHEAD with WILL
	switch cmd {
	case do:
		if opt == optEcho || opt == optSuppressGoAhead {
			_, _ = tc.Conn.Write([]byte{iac, will, opt})
		} else {
			_, _ = tc.Conn.Write([]byte{iac, wont, opt})
		}
	case dont:
		_, _ = tc.Conn.Write([]byte{iac, wont, opt})
	}
}

func (tc *TelnetConn) Write(p []byte) (n int, err error) {
	// Ensure \n is transmitted as \r\n for telnet standard compatibility
	var out bytes.Buffer
	for _, b := range p {
		if b == '\n' {
			out.WriteByte('\r')
		}
		out.WriteByte(b)
	}
	return tc.Conn.Write(out.Bytes())
}