package transport

import (
	"agentdb/internal/model"
	"bufio"
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"io"
	"net"
	"strconv"
	"sync"
	"time"
)

func TLSConfig(p model.Proxy) (*tls.Config, error) {
	if !p.TargetTLS {
		return nil, nil
	}
	cfg := &tls.Config{ServerName: p.TargetHost, MinVersion: tls.VersionTLS12}
	if p.TargetCA != "" {
		pool, err := x509.SystemCertPool()
		if err != nil {
			pool = x509.NewCertPool()
		}
		if !pool.AppendCertsFromPEM([]byte(p.TargetCA)) {
			return nil, errors.New("invalid target CA")
		}
		cfg.RootCAs = pool
	}
	return cfg, nil
}
func Dial(ctx context.Context, p model.Proxy, dial model.DialFunc) (net.Conn, error) {
	c, err := dial(ctx, "tcp", net.JoinHostPort(p.TargetHost, strconv.Itoa(p.TargetPort)))
	if err != nil {
		return nil, err
	}
	cfg, err := TLSConfig(p)
	if err != nil {
		c.Close()
		return nil, err
	}
	if cfg != nil {
		secure := tls.Client(c, cfg)
		if err = secure.HandshakeContext(ctx); err != nil {
			c.Close()
			return nil, err
		}
		c = secure
	}
	return c, nil
}

type BufferedConn struct {
	net.Conn
	Reader *bufio.Reader
}

func (c *BufferedConn) Read(p []byte) (int, error) { return c.Reader.Read(p) }
func Accept(ctx context.Context, raw net.Conn, cfg *tls.Config) (net.Conn, error) {
	raw.SetDeadline(time.Now().Add(15 * time.Second))
	reader := bufio.NewReader(raw)
	first, err := reader.Peek(1)
	if err != nil {
		return nil, err
	}
	var c net.Conn = &BufferedConn{Conn: raw, Reader: reader}
	if first[0] == 22 {
		if cfg == nil {
			return nil, errors.New("gateway TLS certificate missing")
		}
		config := cfg.Clone()
		config.MinVersion = tls.VersionTLS12
		secure := tls.Server(c, config)
		if err = secure.HandshakeContext(ctx); err != nil {
			return nil, err
		}
		c = secure
	}
	return c, nil
}
func Write(w io.Writer, b []byte) error {
	for len(b) > 0 {
		n, err := w.Write(b)
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrShortWrite
		}
		b = b[n:]
	}
	return nil
}

type LockedWriter struct {
	mu     sync.Mutex
	Writer io.Writer
}

func (w *LockedWriter) Write(b []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	err := Write(w.Writer, b)
	if err != nil {
		return 0, err
	}
	return len(b), nil
}