package graceful

import (
	"crypto/tls"
	"errors"
	"net"
	"net/http"
	"sync"
	"time"
)

// TLSListener is a TLS wrapper for net.TCPListener
type TLSListener struct {
	// protects conns
	sync.Mutex

	*net.TCPListener

	config    *tls.Config
	conns     map[net.Conn]struct{}
	keepAlive time.Duration
}

// TLSConn is a net.Conn wrapper
type TLSConn struct {
	net.Conn
	*TLSListener
}

// Close wraps the net.Conn.Close method and removes the conn from the listener map
func (c *TLSConn) Close() error {
	c.RemoveConn(c.Conn)
	return c.Conn.Close()
}

// NewTLS wraps TLSListener around net.TCPListener
func NewTLS(laddr string, config *tls.Config, keepAlive time.Duration) (*TLSListener, error) {
	if config == nil || len(config.Certificates) == 0 {
		return nil, errors.New("NewTLS: no certificates in configuration")
	}
	l, err := net.Listen(network, laddr)
	if err != nil {
		return nil, err
	}
	ltcp, ok := l.(*net.TCPListener)
	if !ok {
		return nil, ErrInterfaceConversion
	}
	return &TLSListener{
		TCPListener: ltcp,
		config:      config,
		conns:       make(map[net.Conn]struct{}),
		keepAlive:   keepAlive,
	}, nil
}

// Accept wraps the net.TCPListener.AcceptTCP method
func (l *TLSListener) Accept() (net.Conn, error) {
	conn, err := l.AcceptTCP()
	if err != nil {
		return conn, err
	}
	if l.keepAlive > 0 {
		conn.SetKeepAlive(true)
		conn.SetKeepAlivePeriod(l.keepAlive)
	}
	tlsconn := &TLSConn{tls.Server(conn, l.config), l}
	l.Lock()
	defer l.Unlock()
	l.conns[tlsconn] = struct{}{}
	return tlsconn, nil
}

// TLSConnState implements HTTP connection hijacking
func (l *TLSListener) TLSConnState(c net.Conn, cs http.ConnState) {
	if cs == http.StateClosed || cs == http.StateHijacked {
		l.RemoveConn(c)
	}
}

// RemoveConn removes a connection from the map
func (l *TLSListener) RemoveConn(c net.Conn) {
	l.Lock()
	defer l.Unlock()
	delete(l.conns, c)
}

// Stop stops the listener and closes all open connections
func (l *TLSListener) Stop() {
	l.Close()
	for c := range l.conns {
		c.Close()
	}
}