//go:build linux

package microvm

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/mdlayher/vsock"
	"tangled.org/core/spindle/observability"
	"tangled.org/core/spindle/quota"
)

type SubstituterUploadBackend interface {
	http.Handler
	Close() error
}

type UploadCacheOptions struct {
	TrustedPublicKeys    []string
	RequireSignedUploads bool
	MaxStagedBytes       int64
	MaxStagedObjects     int
}

type SubstituterUploadProxy struct {
	port uint32

	ln      *vsock.Listener
	server  *http.Server
	backend SubstituterUploadBackend
}

func StartSubstituterUploadProxy(ctx context.Context, cid uint32, uploadURL string, readUpstreams []SubstituterUpstream, stagingDir string, logger *slog.Logger, quotaStore quota.ReservationStore, ownerDID, repoDID string, options ...UploadCacheOptions) (*SubstituterUploadProxy, error) {
	if strings.TrimSpace(uploadURL) == "" {
		return nil, nil
	}
	var opts UploadCacheOptions
	if len(options) > 0 {
		opts = options[0]
	}
	if opts.RequireSignedUploads && len(opts.TrustedPublicKeys) == 0 {
		return nil, errors.New("signed cache uploads require operator-configured trusted public keys")
	}

	if logger == nil {
		logger = slog.Default()
	}
	logger = logger.With("where", "upload_cache_proxy", "cid", cid, "uploadURL", cacheURLForLog(uploadURL))

	backend, err := newSubstituterUploadBackend(uploadURL, operatorSubstituterUpstreams(readUpstreams), stagingDir, logger, opts.TrustedPublicKeys...)
	if err != nil {
		return nil, err
	}

	metrics := observability.GetMetrics(ctx)
	finalBackend := NewStagingWrapper(backend, stagingDir, quotaStore, ownerDID, repoDID, uploadURL, logger, metrics)
	staging := finalBackend.(*stagingWrapper)
	staging.verifier = execStagedNarVerifier{}
	staging.trustedPublicKeys = append([]string(nil), opts.TrustedPublicKeys...)
	staging.requireSignedUploads = opts.RequireSignedUploads
	if nixBackend, ok := backend.(*NixStoreUploadBackend); ok {
		nixBackend.requireSignedUploads = opts.RequireSignedUploads
	}
	if opts.MaxStagedBytes > 0 {
		staging.maxStagedBytes = opts.MaxStagedBytes
	}
	if opts.MaxStagedObjects > 0 {
		staging.maxStagedObjects = opts.MaxStagedObjects
	}

	ln, port, err := listenRandomVsockUploadPort(ctx)
	if err != nil {
		return nil, fmt.Errorf("listen for cache upload proxy: %w", err)
	}

	proxy := &SubstituterUploadProxy{
		port:    port,
		ln:      ln,
		backend: finalBackend,
	}
	proxy.server = &http.Server{
		Handler:           finalBackend,
		Protocols:         cacheProxyProtocols(),
		ReadHeaderTimeout: 30 * time.Second,
	}

	filtered := &cidFilteredVsockListener{
		Listener: ln,
		cid:      cid,
		logger:   logger,
	}
	go func() {
		if err := proxy.server.Serve(filtered); err != nil && !errors.Is(err, http.ErrServerClosed) && !errors.Is(err, net.ErrClosed) {
			logger.Warn("upload cache proxy stopped", "port", port, "error", err)
		}
	}()

	logger.Info("started upload cache proxy", "port", port, "target", cacheURLForLog(uploadURL), "readUpstreams", len(readUpstreams))
	return proxy, nil
}

func newSubstituterUploadBackend(uploadURL string, readUpstreams []SubstituterUpstream, stagingDir string, logger *slog.Logger, trustedPublicKeys ...string) (SubstituterUploadBackend, error) {
	if strings.TrimSpace(uploadURL) == "" {
		return nil, nil
	}

	target, err := url.Parse(uploadURL)
	if err != nil {
		return nil, fmt.Errorf("invalid upload URL %s", cacheURLForLog(uploadURL))
	}

	switch target.Scheme {
	case "http", "https":
		if target.Host == "" {
			return nil, fmt.Errorf("upload URL %s is missing host", cacheURLForLog(uploadURL))
		}
		return newHTTPUploadProxyBackend(target, readUpstreams, logger), nil

	case "ssh", "ssh-ng":
		return newNixStoreUploadBackend(target.String(), stagingDir, readUpstreams, logger, nil, trustedPublicKeys...)

	case "":
		switch uploadURL {
		case "daemon", "local":
			return newNixStoreUploadBackend(uploadURL, stagingDir, readUpstreams, logger, nil, trustedPublicKeys...)
		default:
			return nil, fmt.Errorf("unsupported upload URL %s", cacheURLForLog(uploadURL))
		}

	default:
		return nil, fmt.Errorf("upload URL %s uses unsupported scheme %q", cacheURLForLog(uploadURL), target.Scheme)
	}
}

func (p *SubstituterUploadProxy) Port() uint32 {
	if p == nil {
		return 0
	}
	return p.port
}

func (p *SubstituterUploadProxy) Close() error {
	if p == nil {
		return nil
	}

	var closeErr error
	if p.server != nil {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		closeErr = errors.Join(closeErr, p.server.Shutdown(ctx))
		cancel()
		p.server = nil
	}
	if p.ln != nil {
		closeErr = errors.Join(closeErr, p.ln.Close())
		p.ln = nil
	}
	if p.backend != nil {
		closeErr = errors.Join(closeErr, p.backend.Close())
	}
	return closeErr
}

func listenRandomVsockUploadPort(ctx context.Context) (*vsock.Listener, uint32, error) {
	var lastErr error
	for range 32 {
		port, err := randomVsockPort()
		if err != nil {
			return nil, 0, err
		}
		ln, err := vsock.ListenContextID(vsock.Host, port, nil)
		if err == nil {
			return ln, port, nil
		}
		lastErr = err

		select {
		case <-ctx.Done():
			return nil, 0, ctx.Err()
		default:
		}
	}
	return nil, 0, fmt.Errorf("listen on random vsock upload port: %w", lastErr)
}