package main

import (
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"go/format"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

func main() {
	privateOut := flag.String("private-out", "", "根私钥 PEM 输出路径")
	sdkPublicOut := flag.String("sdk-public-out", "", "SDK 根公钥 Go 源码输出路径")
	flag.Parse()

	if *privateOut == "" || *sdkPublicOut == "" {
		flag.Usage()
		os.Exit(2)
	}
	if sameOutputPath(*privateOut, *sdkPublicOut) {
		fmt.Fprintln(os.Stderr, "根私钥和 SDK 根公钥不能写入同一路径")
		os.Exit(2)
	}

	publicKey, privateKey, err := generateRoot()
	if err != nil {
		fmt.Fprintf(os.Stderr, "生成根密钥失败: %v\n", err)
		os.Exit(1)
	}
	if err := writePrivateKey(*privateOut, privateKey); err != nil {
		fmt.Fprintf(os.Stderr, "写入根私钥失败: %v\n", err)
		os.Exit(1)
	}
	if err := writePublicGo(*sdkPublicOut, publicKey); err != nil {
		if cleanupErr := os.Remove(*privateOut); cleanupErr != nil && !errors.Is(cleanupErr, os.ErrNotExist) {
			err = errors.Join(err, fmt.Errorf("清理本次创建的根私钥失败: %w", cleanupErr))
		}
		fmt.Fprintf(os.Stderr, "写入 SDK 根公钥失败: %v\n", err)
		os.Exit(1)
	}
}

func sameOutputPath(firstPath string, secondPath string) bool {
	firstAbsolute, firstErr := filepath.Abs(firstPath)
	secondAbsolute, secondErr := filepath.Abs(secondPath)
	if firstErr != nil || secondErr != nil {
		return firstPath == secondPath
	}

	if runtime.GOOS == "windows" {
		return strings.EqualFold(firstAbsolute, secondAbsolute)
	}
	return firstAbsolute == secondAbsolute
}

func generateRoot() (ed25519.PublicKey, ed25519.PrivateKey, error) {
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("生成 Ed25519 根密钥: %w", err)
	}
	return publicKey, privateKey, nil
}

func writePrivateKey(filePath string, privateKey ed25519.PrivateKey) error {
	der, err := x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		return fmt.Errorf("编码 PKCS#8 私钥: %w", err)
	}

	pemData := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})
	if pemData == nil {
		return errors.New("编码 PEM 私钥失败")
	}
	return writeNewFile(filePath, pemData, 0o600)
}

func writePublicGo(filePath string, publicKey ed25519.PublicKey) error {
	if len(publicKey) != ed25519.PublicKeySize {
		return fmt.Errorf("根公钥长度为 %d，期望 %d", len(publicKey), ed25519.PublicKeySize)
	}

	var source bytes.Buffer
	source.WriteString("package licence\n\nimport \"crypto/ed25519\"\n\nvar rootPublicKey = [ed25519.PublicKeySize]byte{")
	for index, value := range publicKey {
		if index%8 == 0 {
			source.WriteString("\n\t")
		}
		fmt.Fprintf(&source, "0x%02x, ", value)
	}
	source.WriteString("\n}\n")

	formatted, err := format.Source(source.Bytes())
	if err != nil {
		return fmt.Errorf("格式化 SDK 根公钥源码: %w", err)
	}
	return writeNewFile(filePath, formatted, 0o644)
}

func writeNewFile(filePath string, data []byte, mode os.FileMode) (resultErr error) {
	file, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode)
	if err != nil {
		return err
	}
	closed := false
	completed := false
	defer func() {
		if completed {
			return
		}
		if !closed {
			if closeErr := file.Close(); closeErr != nil {
				resultErr = errors.Join(resultErr, fmt.Errorf("关闭未完成的输出文件: %w", closeErr))
			}
		}
		if removeErr := os.Remove(filePath); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
			resultErr = errors.Join(resultErr, fmt.Errorf("清理未完成的输出文件: %w", removeErr))
		}
	}()

	written, err := file.Write(data)
	if err != nil {
		return fmt.Errorf("写入输出文件: %w", err)
	}
	if written != len(data) {
		return fmt.Errorf("写入输出文件不完整: 已写入 %d 字节，期望 %d 字节", written, len(data))
	}
	if err := file.Close(); err != nil {
		closed = true
		return fmt.Errorf("关闭输出文件: %w", err)
	}
	closed = true
	completed = true
	return nil
}