//go:build windows

package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/netip"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/tailscale/wireguard-go/tun"
	"golang.org/x/sys/windows"
	"tailscale.com/cmd/mirage/internal/subprocessenv"
	"tailscale.com/control/controlclient"
	"tailscale.com/drive/driveimpl"
	"tailscale.com/health"
	"tailscale.com/ipn"
	"tailscale.com/ipn/ipnlocal"
	"tailscale.com/ipn/ipnserver"
	"tailscale.com/ipn/store"
	"tailscale.com/logpolicy"
	"tailscale.com/net/dns"
	"tailscale.com/net/dnsfallback"
	"tailscale.com/net/netmon"
	"tailscale.com/net/netns"
	"tailscale.com/net/tsdial"
	"tailscale.com/net/tstun"
	"tailscale.com/paths"
	"tailscale.com/safesocket"
	"tailscale.com/syncs"
	"tailscale.com/tsd"
	"tailscale.com/types/logger"
	"tailscale.com/types/logid"
	"tailscale.com/util/backoff"
	"tailscale.com/util/miragebrand"
	"tailscale.com/util/winutil"
	"tailscale.com/wgengine"
	"tailscale.com/wgengine/netstack"
	"tailscale.com/wgengine/router"
)

var runtimeLogPolicy *logpolicy.Policy

func init() {
	tstunNew = tstunNewWithWindowsRetries
}

type daemonOptions struct {
	TunName    string
	ListenPort uint16
	StatePath  string
	StateDir   string
	SocketPath string
}

func defaultDaemonOptions() daemonOptions {
	return daemonOptions{
		TunName:    miragebrand.WindowsTunName(),
		ListenPort: 41641,
		StatePath:  paths.DefaultTailscaledStateFile(),
		StateDir:   paths.DefaultTailscaledStateDir(),
		SocketPath: paths.DefaultTailscaledSocket(),
	}
}

func statePathOrDefault(opts daemonOptions) string {
	if opts.StatePath != "" {
		return opts.StatePath
	}
	if opts.StateDir != "" {
		return filepath.Join(opts.StateDir, "tailscaled.state")
	}
	return ""
}

type serverOptions struct {
	VarRoot    string
	LoginFlags controlclient.LoginFlags
}

func ipnServerOpts(opts daemonOptions) (o serverOptions) {
	o.VarRoot = opts.StateDir
	if o.VarRoot == "" && filepath.IsAbs(opts.StatePath) {
		if dir := filepath.Dir(opts.StatePath); strings.EqualFold(filepath.Base(dir), "mirage") {
			o.VarRoot = dir
		}
	}
	if strings.HasPrefix(statePathOrDefault(opts), "mem:") {
		o.LoginFlags = controlclient.LoginEphemeral
	}
	return o
}

func runMirageSubprocess() error {
	log.SetFlags(0)
	log.Printf("mirage subprocess entrypoint starting")

	exePath, err := os.Executable()
	if err != nil {
		return fmt.Errorf("os.Executable: %w", err)
	}
	if err := subprocessenv.Configure(
		exePath,
		windows.SetDllDirectory,
		func(path string) error {
			_, err := windows.LoadDLL(path)
			return err
		},
		log.Printf,
	); err != nil {
		return err
	}
	log.Printf("Program starting: %#v", os.Args)

	logID := logid.PublicID{}
	if len(os.Args) > 2 {
		parsed, err := logid.ParsePublicID(os.Args[2])
		if err != nil {
			log.Printf("mirage subprocess logid parse failed for %q: %v", os.Args[2], err)
		} else {
			logID = parsed
			log.Printf("mirage subprocess logid: %s", logID)
		}
	}

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

	go func() {
		buf := make([]byte, 16)
		for {
			_, err := os.Stdin.Read(buf)
			if errors.Is(err, os.ErrClosed) || errors.Is(err, io.EOF) {
				cancel()
				return
			}
			if err != nil {
				log.Fatalf("stdin read error: %v", err)
			}
		}
	}()

	sys := tsd.NewSystem()
	logf := logger.Logf(log.Printf)
	log.Printf("mirage subprocess runtime initialized; entering backend bootstrap")

	return startMirageIPNServer(ctx, logf, logID, sys, defaultDaemonOptions(), func() error {
		logf("subproc stage: initializing netmon")
		netMon, err := netmon.New(sys.Bus.Get(), logf)
		if err != nil {
			return fmt.Errorf("netmon.New: %w", err)
		}
		sys.Set(netMon)
		logf("subproc stage: initializing drive filesystem")
		sys.Set(driveimpl.NewFileSystemForRemote(logf))
		return nil
	})
}

func tstunNewWithWindowsRetries(logf logger.Logf, tunName string) (_ tun.Device, devName string, _ error) {
	bo := backoff.NewBackoff("tstunNew", logf, 10*time.Second)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()
	for {
		dev, devName, err := tstun.New(logf, tunName)
		if err == nil {
			return dev, devName, nil
		}
		if errors.Is(err, windows.ERROR_DEVICE_NOT_AVAILABLE) || windowsUptime() < 10*time.Minute {
			winutil.LogSvcState(logf, "NetSetupSvc")
		}
		bo.BackOff(ctx, err)
		if ctx.Err() != nil {
			return nil, "", ctx.Err()
		}
	}
}

var (
	kernel32Daemon           = windows.NewLazySystemDLL("kernel32.dll")
	getTickCount64ProcDaemon = kernel32Daemon.NewProc("GetTickCount64")
)

func windowsUptime() time.Duration {
	r, _, _ := getTickCount64ProcDaemon.Call()
	return time.Duration(int64(r)) * time.Millisecond
}

func startMirageIPNServer(ctx context.Context, logf logger.Logf, logID logid.PublicID, sys *tsd.System, opts daemonOptions, prepareBackend func() error) error {
	logf("subproc stage: creating LocalAPI listener on %q", opts.SocketPath)
	ln, err := safesocket.Listen(opts.SocketPath)
	if err != nil {
		return fmt.Errorf("safesocket.Listen: %w", err)
	}
	logf("subproc stage: LocalAPI listener ready on %q", opts.SocketPath)

	if prepareBackend != nil {
		logf("subproc stage: preparing backend prerequisites")
		if err := prepareBackend(); err != nil {
			return fmt.Errorf("prepare backend prerequisites: %w", err)
		}
		logf("subproc stage: backend prerequisites ready")
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM)

	wgEngineCreated := make(chan struct{})
	go func() {
		var wgEngineClosed <-chan struct{}
		waitEngine := wgEngineCreated
		for {
			select {
			case <-interrupt:
				cancel()
				return
			case <-wgEngineClosed:
				cancel()
				return
			case <-waitEngine:
				wgEngineClosed = sys.Engine.Get().Done()
				waitEngine = nil
			case <-ctx.Done():
				return
			}
		}
	}()

	logf("subproc stage: constructing LocalAPI server")
	srv := ipnserver.New(logf, logID, sys.Bus.Get(), sys.NetMon.Get())
	var lbErr syncs.AtomicValue[error]

	go func() {
		logf("subproc stage: initializing LocalBackend")
		lb, err := getLocalBackend(ctx, logf, logID, sys, opts)
		if err != nil {
			logf("subproc stage: LocalBackend initialization failed: %v", err)
			lbErr.Store(err)
			cancel()
			return
		}
		if lb.Prefs().Valid() {
			if err := lb.Start(ipn.Options{}); err != nil {
				logf("subproc stage: LocalBackend.Start failed: %v", err)
				lb.Shutdown()
				lbErr.Store(err)
				cancel()
				return
			}
		}
		srv.SetLocalBackend(lb)
		logf("subproc stage: LocalBackend attached to LocalAPI server")
		close(wgEngineCreated)
	}()

	logf("subproc stage: running LocalAPI server")
	err = srv.Run(ctx, ln)
	if err != nil && lbErr.Load() != nil {
		return fmt.Errorf("getLocalBackend error: %v", lbErr.Load())
	}
	if err != nil && !errors.Is(err, context.Canceled) {
		return fmt.Errorf("ipnserver.Run: %w", err)
	}
	return nil
}

func getLocalBackend(ctx context.Context, logf logger.Logf, logID logid.PublicID, sys *tsd.System, opts daemonOptions) (_ *ipnlocal.LocalBackend, retErr error) {
	if runtimeLogPolicy != nil {
		runtimeLogPolicy.Logtail.SetNetMon(sys.NetMon.Get())
	}

	logf("backend stage: constructing dialer")
	dialer := &tsdial.Dialer{Logf: logf}
	dialer.SetBus(sys.Bus.Get())
	sys.Set(dialer)

	logf("backend stage: creating engine")
	onlyNetstack, err := createEngine(logf, sys, opts)
	if err != nil {
		return nil, fmt.Errorf("createEngine: %w", err)
	}

	logf("backend stage: creating netstack")
	ns, err := newNetstackImpl(logf, sys, onlyNetstack)
	if err != nil {
		return nil, fmt.Errorf("newNetstack: %w", err)
	}

	logf("backend stage: loading state store")
	storePath := statePathOrDefault(opts)
	if storePath == "" {
		return nil, errors.New("mirage state path is empty")
	}
	stateStore, err := store.New(logf, storePath)
	if err != nil {
		ht, ok := sys.HealthTracker.GetOK()
		if !ok {
			return nil, fmt.Errorf("store.New: %w", err)
		}
		logf("store.New failed: %v; starting with in-memory store with a health warning", err)
		stateStore = new(memStoreShim)
		ht.SetUnhealthy(ipn.StateStoreHealth, health.Args{health.ArgError: err.Error()})
	}
	sys.Set(stateStore)

	if tunDev, ok := sys.Tun.GetOK(); ok {
		logf("backend stage: starting tun device")
		tunDev.Start()
	}

	logf("backend stage: creating LocalBackend")
	lb, err := ipnlocal.NewLocalBackend(logf, logID, sys, ipnServerOpts(opts).LoginFlags)
	if err != nil {
		return nil, fmt.Errorf("ipnlocal.NewLocalBackend: %w", err)
	}
	lb.SetVarRoot(ipnServerOpts(opts).VarRoot)
	if runtimeLogPolicy != nil {
		lb.SetLogFlusher(runtimeLogPolicy.Logtail.StartFlush)
	}
	if root := lb.TailscaleVarRoot(); root != "" {
		dnsfallback.SetCachePath(filepath.Join(root, "derpmap.cached.json"), logf)
	}
	logf("backend stage: starting netstack")
	if err := ns.Start(lb); err != nil {
		return nil, fmt.Errorf("netstack.Start: %w", err)
	}
	return lb, nil
}

func newNetstackImpl(logf logger.Logf, sys *tsd.System, onlyNetstack bool) (tsd.NetstackImpl, error) {
	ns, err := netstack.Create(
		logf,
		sys.Tun.Get(),
		sys.Engine.Get(),
		sys.MagicSock.Get(),
		sys.Dialer.Get(),
		sys.DNSManager.Get(),
		sys.ProxyMapper(),
	)
	if err != nil {
		return nil, err
	}
	sys.Set(ns)
	ns.ProcessLocalIPs = onlyNetstack
	ns.ProcessSubnets = onlyNetstack || handleSubnetsInNetstack()

	dialer := sys.Dialer.Get()
	if onlyNetstack {
		engine := sys.Engine.Get()
		dialer.UseNetstackForIP = func(ip netip.Addr) bool {
			_, ok := engine.PeerForIP(ip)
			return ok
		}
		dialer.NetstackDialTCP = func(ctx context.Context, dst netip.AddrPort) (net.Conn, error) {
			conn, err := ns.DialContextTCP(ctx, dst)
			if err != nil {
				return nil, err
			}
			return conn, nil
		}
		dialer.NetstackDialUDP = func(ctx context.Context, dst netip.AddrPort) (net.Conn, error) {
			conn, err := ns.DialContextUDP(ctx, dst)
			if err != nil {
				return nil, err
			}
			return conn, nil
		}
	}
	return ns, nil
}

func handleSubnetsInNetstack() bool {
	return true
}

var tstunNew = tstun.New

func createEngine(logf logger.Logf, sys *tsd.System, opts daemonOptions) (onlyNetstack bool, err error) {
	if opts.TunName == "" {
		return false, errors.New("no tunnel name configured")
	}
	conf := wgengine.Config{
		ListenPort:    opts.ListenPort,
		NetMon:        sys.NetMon.Get(),
		HealthTracker: sys.HealthTracker.Get(),
		Metrics:       sys.UserMetricsRegistry(),
		Dialer:        sys.Dialer.Get(),
		SetSubsystem:  sys.Set,
		ControlKnobs:  sys.ControlKnobs(),
		EventBus:      sys.Bus.Get(),
	}

	onlyNetstack = opts.TunName == "userspace-networking"
	netns.SetEnabled(!onlyNetstack)
	if onlyNetstack {
		e, err := wgengine.NewUserspaceEngine(logf, conf)
		if err != nil {
			return false, err
		}
		sys.Set(wgengine.NewWatchdog(e))
		sys.NetstackRouter.Set(true)
		return true, nil
	}

	dev, devName, err := tstunNew(logf, opts.TunName)
	if err != nil {
		tstun.Diagnose(logf, opts.TunName, err)
		return false, fmt.Errorf("tstun.New(%q): %w", opts.TunName, err)
	}
	conf.Tun = dev

	routerDev, err := router.New(logf, dev, sys.NetMon.Get(), sys.HealthTracker.Get(), sys.Bus.Get())
	if err != nil {
		dev.Close()
		return false, fmt.Errorf("creating router: %w", err)
	}
	dnsMgr, err := dns.NewOSConfigurator(logf, sys.HealthTracker.Get(), sys.Bus.Get(), sys.PolicyClientOrDefault(), sys.ControlKnobs(), devName)
	if err != nil {
		dev.Close()
		routerDev.Close()
		return false, fmt.Errorf("dns.NewOSConfigurator: %w", err)
	}
	conf.Router = routerDev
	conf.DNS = dnsMgr
	sys.Set(conf.Router)

	engine, err := wgengine.NewUserspaceEngine(logf, conf)
	if err != nil {
		return false, err
	}
	sys.Set(wgengine.NewWatchdog(engine))
	sys.NetstackRouter.Set(handleSubnetsInNetstack())
	return false, nil
}

type memStoreShim struct{}

func (*memStoreShim) ReadState(ipn.StateKey) ([]byte, error) { return nil, ipn.ErrStateNotExist }
func (*memStoreShim) WriteState(ipn.StateKey, []byte) error  { return nil }
