//go:build windows

package divert

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

// DriverPaths records where the WinDivert binaries landed after install.
type DriverPaths struct {
	SysPath string // e.g. C:\ProgramData\Drover\windivert\WinDivert64.sys
	DllPath string
}

// InstallDriver extracts the embedded WinDivert.sys + WinDivert.dll
// into %PROGRAMDATA%\Drover\windivert\ and SHA256-verifies them.
//
// On second and subsequent runs, if the existing files already match
// the embedded SHAs, the function is a no-op and just returns paths.
//
// Errors:
//   - ARM64 architecture (WinDivert doesn't support it)
//   - %PROGRAMDATA% not set or not writable
//   - SHA256 mismatch after write (driver corrupted on disk)
func InstallDriver() (*DriverPaths, error) {
	if runtime.GOARCH == "arm64" {
		return nil, fmt.Errorf("Drover requires x86-64 Windows; ARM64 is not supported (WinDivert does not ship an ARM64 driver)")
	}
	pd := os.Getenv("ProgramData")
	if pd == "" {
		return nil, fmt.Errorf("ProgramData environment variable is not set")
	}
	dst := filepath.Join(pd, "Drover", "windivert")
	return installDriverInto(dst)
}

func installDriverInto(dst string) (*DriverPaths, error) {
	if runtime.GOARCH == "arm64" {
		return nil, fmt.Errorf("Drover requires x86-64 Windows; ARM64 is not supported")
	}
	if err := os.MkdirAll(dst, 0755); err != nil {
		return nil, fmt.Errorf("create %s: %w", dst, err)
	}
	sysPath := filepath.Join(dst, "WinDivert64.sys")
	dllPath := filepath.Join(dst, "WinDivert.dll")

	if err := writeIfDifferent(sysPath, winDivertSys, WinDivertSysSHA256); err != nil {
		return nil, fmt.Errorf("install WinDivert64.sys: %w", err)
	}
	if err := writeIfDifferent(dllPath, winDivertDll, WinDivertDllSHA256); err != nil {
		return nil, fmt.Errorf("install WinDivert.dll: %w", err)
	}
	return &DriverPaths{SysPath: sysPath, DllPath: dllPath}, nil
}

// writeIfDifferent compares the existing file's SHA256 to the expected
// hash; if it matches, no-op. Otherwise overwrite atomically and verify
// the resulting on-disk SHA matches expected.
func writeIfDifferent(path string, content []byte, expectedSHA string) error {
	if existing, err := os.ReadFile(path); err == nil {
		if strings.EqualFold(sha256Hex(existing), expectedSHA) {
			return nil // already up to date
		}
	}
	tmp := path + ".new"
	if err := os.WriteFile(tmp, content, 0644); err != nil {
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	// Verify after write — guards against AV-on-write tampering.
	got, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	if !strings.EqualFold(sha256Hex(got), expectedSHA) {
		return fmt.Errorf("SHA256 mismatch after write at %s; antivirus may have tampered with the file. Add %%PROGRAMDATA%%\\Drover\\ to your AV exclusions and restart Drover", path)
	}
	return nil
}

func sha256Hex(b []byte) string {
	h := sha256.Sum256(b)
	return hex.EncodeToString(h[:])
}