package vless

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/xtls/reality"

	"go.bigb.es/shroud/internal/store"
)

// Config holds the VLESS+REALITY server configuration.
type Config struct {
	ListenAddr  string
	PrivateKey  []byte   // 32-byte raw x25519 private key
	PublicKey   []byte   // 32-byte raw x25519 public key (for share links)
	ShortID     [8]byte
	ServerNames []string
	Dest        string
	Show        bool
}

// Server is a VLESS proxy server with REALITY obfuscation.
type Server struct {
	cfg      Config
	listener net.Listener
	logger   *slog.Logger

	mu     sync.RWMutex
	uuids  map[[16]byte]string // UUID → access key ID
	cancel context.CancelFunc
}

// New creates a new VLESS+REALITY server.
func New(cfg Config, logger *slog.Logger) *Server {
	return &Server{
		cfg:    cfg,
		logger: logger.With("component", "vless"),
		uuids:  make(map[[16]byte]string),
	}
}

// Start initializes the REALITY listener and begins accepting connections.
func (s *Server) Start() error {
	shortIds := map[[8]byte]bool{s.cfg.ShortID: true}

	serverNames := make(map[string]bool, len(s.cfg.ServerNames))
	for _, name := range s.cfg.ServerNames {
		serverNames[name] = true
	}

	realityCfg := &reality.Config{
		PrivateKey:  s.cfg.PrivateKey,
		ShortIds:    shortIds,
		ServerNames: serverNames,
		Dest:        s.cfg.Dest,
		Show:        s.cfg.Show,
	}

	ln, err := reality.Listen("tcp", s.cfg.ListenAddr, realityCfg)
	if err != nil {
		return fmt.Errorf("reality listen: %w", err)
	}
	s.listener = ln

	ctx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel

	go s.acceptLoop(ctx)

	return nil
}

// SyncKeys rebuilds the UUID authentication map from the given access keys.
func (s *Server) SyncKeys(keys []store.AccessKey) error {
	newUUIDs := make(map[[16]byte]string)
	for _, k := range keys {
		if k.VLESS == nil {
			continue
		}
		parsed, err := uuid.Parse(k.VLESS.UUID)
		if err != nil {
			s.logger.Warn("Skipping VLESS key with invalid UUID.", "keyID", k.ID, "err", err)
			continue
		}
		newUUIDs[parsed] = k.ID
	}

	s.mu.Lock()
	s.uuids = newUUIDs
	s.mu.Unlock()

	s.logger.Info("VLESS keys synced.", "count", len(newUUIDs))
	return nil
}

// Stop shuts down the server.
func (s *Server) Stop() {
	if s.cancel != nil {
		s.cancel()
	}
	if s.listener != nil {
		s.listener.Close()
	}
	s.logger.Info("VLESS server stopped.")
}

func (s *Server) acceptLoop(ctx context.Context) {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				s.logger.Error("VLESS accept error.", "err", err)
				continue
			}
		}
		go s.handleConn(conn)
	}
}

func (s *Server) handleConn(conn net.Conn) {
	defer conn.Close()

	// Deadline for header parsing.
	conn.SetReadDeadline(time.Now().Add(10 * time.Second))

	req, err := ParseRequest(conn)
	if err != nil {
		s.logger.Debug("Failed to parse VLESS request.", "err", err)
		return
	}

	// Clear read deadline for relay.
	conn.SetReadDeadline(time.Time{})

	// Authenticate UUID.
	s.mu.RLock()
	keyID, ok := s.uuids[req.UUID]
	s.mu.RUnlock()

	if !ok {
		s.logger.Debug("Unknown VLESS UUID, closing connection.")
		return
	}

	target := req.Target()

	switch req.Command {
	case CmdTCP:
		s.logger.Debug("VLESS TCP relay.", "keyID", keyID, "target", target)
		if err := relayTCP(conn, target); err != nil {
			s.logger.Debug("VLESS TCP relay ended.", "keyID", keyID, "err", err)
		}

	case CmdUDP:
		s.logger.Debug("VLESS UDP relay.", "keyID", keyID, "target", target)
		if err := relayUDP(conn, target); err != nil {
			s.logger.Debug("VLESS UDP relay ended.", "keyID", keyID, "err", err)
		}
	}
}