package shadow

import (
	"crypto/tls"
	"github.com/xtaci/smux"
	"net"
	"sync"
	"time"
)

type TLSBridge struct {
	session        *smux.Session
	locker         sync.Mutex
	serverAddress  string
	fakeAddressSNI string
}

func NewTLSBridge(serverAddress string, fakeAddressSNI string) *TLSBridge {
	t := &TLSBridge{
		session:        nil,
		locker:         sync.Mutex{},
		serverAddress:  serverAddress,
		fakeAddressSNI: fakeAddressSNI,
	}
	return t
}

func (t *TLSBridge) dial() error {
	if t.session != nil {
		t.session.Close()
	}
	dial, err := tls.DialWithDialer(&net.Dialer{
		Timeout: time.Second * 5,
	}, "tcp", t.serverAddress, &tls.Config{
		ServerName: t.fakeAddressSNI,
	})
	if err != nil {
		return err
	}
	err = dial.Handshake()
	if err != nil {
		return err
	}
	session, err := smux.Client(dial.NetConn(), nil)
	if err != nil {
		return err
	}
	//force openStream to prevent first connection problem
	session.OpenStream()
	t.session = session
	return nil
}

func (t *TLSBridge) GetStream() *smux.Stream {
	t.locker.Lock()
	defer t.locker.Unlock()

	if t.session == nil {
		err := t.dial()
		if err != nil {
			return nil
		}
	}

	openStream, err := t.session.OpenStream()
	if err != nil {
		t.session.Close()
		t.session = nil
	}
	return openStream

}