package ssh

import (
	"net"
	"time"

	"github.com/fujiwara/shapeio"
	"github.com/google/uuid"
)

type throttledConn struct {
	ID            uuid.UUID
	conn          net.Conn
	writer        *shapeio.Writer
	reader        *shapeio.Reader
	CloseCallback func(c *throttledConn)
}

func newThrottledConn(conn net.Conn) *throttledConn {
	id := uuid.Must(uuid.NewRandom())
	return &throttledConn{
		ID:     id,
		conn:   conn,
		writer: shapeio.NewWriter(conn),
		reader: shapeio.NewReader(conn),
	}
}

func (sc *throttledConn) SetSpeed(bytesPerSec float64) {
	sc.writer.SetRateLimit(bytesPerSec)
	sc.reader.SetRateLimit(bytesPerSec)
}

func (sc *throttledConn) Read(b []byte) (n int, err error) {
	return sc.reader.Read(b)
}

func (sc *throttledConn) Write(b []byte) (n int, err error) {
	return sc.writer.Write(b)
}

func (sc *throttledConn) Close() error {
	if sc.CloseCallback != nil {
		sc.CloseCallback(sc)
	}
	return sc.conn.Close()
}

func (sc *throttledConn) LocalAddr() net.Addr {
	return sc.conn.LocalAddr()
}

func (sc *throttledConn) RemoteAddr() net.Addr {
	return sc.conn.RemoteAddr()
}

func (sc *throttledConn) SetDeadline(t time.Time) error {
	return sc.conn.SetDeadline(t)
}

func (sc *throttledConn) SetReadDeadline(t time.Time) error {
	return sc.conn.SetReadDeadline(t)
}

func (sc *throttledConn) SetWriteDeadline(t time.Time) error {
	return sc.conn.SetWriteDeadline(t)
}