package net

import (
	"bufio"
	"bytes"
	"io"
	"net"
	"strings"
)

type Conn interface {
	ReadMessage() (string, error)
	io.Writer
	io.Closer
	SetEcho(bool) error
}

type tcpConn struct {
	conn    net.Conn
	reader  *bufio.Reader
	echoOff bool
}

func newTCPConn(conn net.Conn) *tcpConn {
	return &tcpConn{
		conn:   conn,
		reader: bufio.NewReaderSize(conn, maxLineLength*4),
	}
}

// maxLineLength bounds how much of an incoming line is accepted. Longer
// lines are drained from the socket and truncated so a client cannot grow
// the read buffer without limit.
const maxLineLength = 1024

func (t *tcpConn) ReadMessage() (string, error) {
	line, err := t.reader.ReadSlice('\n')
	if err == bufio.ErrBufferFull {
		// Copy the bounded prefix out of the reusable read buffer BEFORE
		// draining the rest of the oversized line; further reads overwrite the
		// buffer the slice points into.
		prefixLen := len(line)
		if prefixLen > maxLineLength {
			prefixLen = maxLineLength
		}
		prefix := make([]byte, prefixLen)
		copy(prefix, line[:prefixLen])
		for {
			_, err = t.reader.ReadSlice('\n')
			if err == nil {
				break
			}
			if err != bufio.ErrBufferFull {
				return "", err
			}
		}
		line = prefix
	} else if err != nil {
		return "", err
	}
	text := strings.TrimRight(string(line), "\r\n")
	text = stripIAC(text)
	if t.echoOff {
		t.conn.Write([]byte("\r\n"))
	}
	return text, nil
}

func (t *tcpConn) Write(b []byte) (int, error) {
	return t.conn.Write(b)
}

func (t *tcpConn) Close() error {
	return t.conn.Close()
}

func (t *tcpConn) SetEcho(on bool) error {
	t.echoOff = !on
	if on {
		return t.writeTelnetCmd(wont, echo)
	}
	return t.writeTelnetCmd(will, echo)
}

const (
	iac  = 255
	will = 251
	wont = 252
	do   = 253
	dont = 254
	sb   = 250
	se   = 240
	echo = 1
)

func (t *tcpConn) writeTelnetCmd(cmd, opt byte) error {
	_, err := t.conn.Write([]byte{iac, cmd, opt})
	return err
}

func stripIAC(s string) string {
	b := []byte(s)
	var out []byte
	i := 0
	iacSE := []byte{iac, se}
	for i < len(b) {
		if b[i] == iac && i+2 < len(b) {
			if b[i+1] == sb {
				end := bytes.Index(b[i+2:], iacSE)
				if end >= 0 {
					i += end + 4
					continue
				}
			}
			i += 3
			continue
		}
		if b[i] == iac {
			i++
			continue
		}
		out = append(out, b[i])
		i++
	}
	return string(out)
}