package peer

import (
	"errors"
	"io"
	"log/slog"
	"net"
	"sync"

	"github.com/shadowsocks/go-shadowsocks2/core"
)

const (
	tcpNetwork              = "tcp"
	bufferSize              = 32 * 1024                // 32KB
	AHEAD_CHACHA20_POLY1305 = "AEAD_CHACHA20_POLY1305" // Default method cipher
)

// Ошибки подключения
var (
	ErrNotConnected = errors.New("peer is not connected")
)

type Peer struct {
	pubKey        string
	connEncrypted net.Conn
	mu            sync.Mutex
	active        bool
}

func NewPeer(pubKey string) *Peer {
	return &Peer{
		pubKey: pubKey,
		active: false,
	}
}

func (p *Peer) Connect(endpoint string) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.active {
		return nil // Уже подключен
	}

	conn, err := net.Dial(tcpNetwork, endpoint)
	if err != nil {
		return err
	}

	cipher, err := core.PickCipher(AHEAD_CHACHA20_POLY1305, []byte(p.pubKey), "")
	if err != nil {
		conn.Close()
		return err
	}

	p.connEncrypted = cipher.StreamConn(conn)
	p.active = true
	return nil
}

func (p *Peer) Disconnect() error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if !p.active {
		return nil
	}

	err := p.connEncrypted.Close()
	p.active = false
	return err
}

func (p *Peer) Start(localClient net.Conn) error {
	if !p.active {
		return ErrNotConnected
	}

	var wg sync.WaitGroup
	wg.Add(2)

	// Логика проксирования в обе стороны
	go p.pipeData(localClient, p.connEncrypted, &wg)
	go p.pipeData(p.connEncrypted, localClient, &wg)

	wg.Wait()
	return nil
}

func (p *Peer) pipeData(src, dst net.Conn, wg *sync.WaitGroup) {
	defer wg.Done()

	buf := make([]byte, bufferSize)
	_, err := io.CopyBuffer(dst, src, buf)
	if err != nil {
		slog.Warn("pipe data error", slog.String("message", err.Error()))
	}

	// Закрываем соединения при завершении
	src.Close()
	dst.Close()
}