package cli

import (
	"bytes"
	"io"
	"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
	state          byte
	optionCommand  byte
	lastCR         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) {
	if len(p) == 0 {
		return 0, nil
	}
	buf := make([]byte, len(p))
	for {
		rawN, readErr := tc.Conn.Read(buf)
		var clean bytes.Buffer
		for _, b := range buf[:rawN] {
			switch tc.state {
			case 1: // IAC
				tc.state = 0
				switch b {
				case iac:
					if !tc.subnegotiation {
						clean.WriteByte(iac)
					}
				case will, wont, do, dont:
					tc.optionCommand = b
					tc.state = 2
				case sb:
					tc.subnegotiation = true
				case se:
					tc.subnegotiation = false
				}
				continue
			case 2:
				tc.handleOption(tc.optionCommand, b)
				tc.state = 0
				continue
			}
			if b == iac {
				tc.state = 1
				continue
			}
			if tc.subnegotiation {
				continue
			}
			if tc.lastCR {
				tc.lastCR = false
				if b == '\n' || b == 0 {
					continue
				}
			}
			if b == '\r' {
				clean.WriteByte('\n')
				tc.lastCR = true
			} else {
				clean.WriteByte(b)
			}
		}
		if clean.Len() > 0 {
			return copy(p, clean.Bytes()), readErr
		}
		if readErr != nil {
			return 0, readErr
		}
	}
}

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)
	}
	written, err := tc.Conn.Write(out.Bytes())
	if err != nil {
		return 0, err
	}
	if written != out.Len() {
		return 0, io.ErrShortWrite
	}
	return len(p), nil
}