package cmdline

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
	"time"

	"github.com/spf13/cobra"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/http/auth"
)

func init() {
	certCmd := cobra.Command{
		Use: "cert",
	}
	RootCmd.AddCommand(&certCmd)

	certRoot := cobra.Command{
		Use:  "root [outputDir]",
		Args: cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			certRoot(args[0])
		},
	}
	certCmd.AddCommand(&certRoot)

	var certFilePath string
	var keyFilePath string

	certServer := cobra.Command{
		Use:  "server [outputDir]",
		Args: cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			certServer(certFilePath, keyFilePath, args[0])
		},
	}
	certServer.Flags().StringVar(&certFilePath, "cert", "", "CA certificate file path")
	certServer.Flags().StringVar(&keyFilePath, "key", "", "CA key file path")
	certCmd.AddCommand(&certServer)

	certClient := cobra.Command{
		Use:  "client [outputDir]",
		Args: cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			certClient(certFilePath, keyFilePath, args[0])
		},
	}
	certClient.Flags().StringVar(&certFilePath, "cert", "", "CA certificate file path")
	certClient.Flags().StringVar(&keyFilePath, "key", "", "CA key file path")
	certCmd.AddCommand(&certClient)
}

func certRoot(output string) {
	caPriv, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

	// 创建 CA 证书模板
	caTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization: []string{"JCS"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(10, 0, 0), // 有效期10年
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}

	// 自签名 CA 证书
	caCertDER, _ := x509.CreateCertificate(rand.Reader, caTemplate, caTemplate, &caPriv.PublicKey, caPriv)

	// 保存 CA 证书和私钥
	writePem(filepath.Join(output, "ca_cert.pem"), "CERTIFICATE", caCertDER)

	privDER, _ := x509.MarshalECPrivateKey(caPriv)
	writePem(filepath.Join(output, "ca_key.pem"), "EC PRIVATE KEY", privDER)
	fmt.Println("CA certificate and key saved to", output)
}

func certServer(certFile string, keyFile string, output string) {
	// 读取 CA 证书和私钥
	caCertPEM, err := os.ReadFile(certFile)
	if err != nil {
		fmt.Println("Failed to read CA certificate:", err)
		return
	}
	caKeyPEM, err := os.ReadFile(keyFile)
	if err != nil {
		fmt.Println("Failed to read CA key:", err)
		return
	}
	caCertPEMBlock, _ := pem.Decode(caCertPEM)
	if caCertPEMBlock == nil {
		fmt.Println("Failed to decode CA certificate")
		return
	}
	caKeyPEMBlock, _ := pem.Decode(caKeyPEM)
	if caKeyPEMBlock == nil {
		fmt.Println("Failed to decode CA key")
		return
	}

	caCert, err := x509.ParseCertificate(caCertPEMBlock.Bytes)
	if err != nil {
		fmt.Println("Failed to parse CA certificate:", err)
		return
	}

	caKey, err := x509.ParseECPrivateKey(caKeyPEMBlock.Bytes)
	if err != nil {
		fmt.Println("Failed to parse CA key:", err)
		return
	}

	// 生成服务端私钥
	serverPriv, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

	// 服务端证书模板
	serverTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		Subject: pkix.Name{
			CommonName: "localhost",
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(1, 0, 0), // 有效期1年
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	// 添加主机名/IP 到证书
	serverTemplate.DNSNames = []string{auth.ClientInternalSNI}

	// 用 CA 签发服务端证书
	serverCertDER, _ := x509.CreateCertificate(rand.Reader, serverTemplate, caCert, &serverPriv.PublicKey, caKey)

	// 保存服务端证书和私钥
	writePem(filepath.Join(output, "server_cert.pem"), "CERTIFICATE", serverCertDER)

	privPem, _ := x509.MarshalECPrivateKey(serverPriv)
	writePem(filepath.Join(output, "server_key.pem"), "EC PRIVATE KEY", privPem)
	fmt.Println("Server certificate and key saved to", output)
}

func certClient(certFile string, keyFile string, output string) {
	// 读取 CA 证书和私钥
	caCertPEM, err := os.ReadFile(certFile)
	if err != nil {
		fmt.Println("Failed to read CA certificate:", err)
		return
	}
	caKeyPEM, err := os.ReadFile(keyFile)
	if err != nil {
		fmt.Println("Failed to read CA key:", err)
		return
	}
	caCertPEMBlock, _ := pem.Decode(caCertPEM)
	if caCertPEMBlock == nil {
		fmt.Println("Failed to decode CA certificate")
		return
	}
	caKeyPEMBlock, _ := pem.Decode(caKeyPEM)
	if caKeyPEMBlock == nil {
		fmt.Println("Failed to decode CA key")
		return
	}

	caCert, err := x509.ParseCertificate(caCertPEMBlock.Bytes)
	if err != nil {
		fmt.Println("Failed to parse CA certificate:", err)
		return
	}

	caKey, err := x509.ParseECPrivateKey(caKeyPEMBlock.Bytes)
	if err != nil {
		fmt.Println("Failed to parse CA key:", err)
		return
	}

	// 生成客户端私钥
	clientPriv, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

	// 客户端证书模板
	clientTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(3),
		Subject: pkix.Name{
			CommonName: "client",
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(1, 0, 0), // 有效期1年
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	// 用 CA 签发客户端证书
	clientCertDER, _ := x509.CreateCertificate(rand.Reader, clientTemplate, caCert, &clientPriv.PublicKey, caKey)

	// 保存客户端证书和私钥
	writePem(filepath.Join(output, "client_cert.pem"), "CERTIFICATE", clientCertDER)

	privPem, _ := x509.MarshalECPrivateKey(clientPriv)
	writePem(filepath.Join(output, "client_key.pem"), "EC PRIVATE KEY", privPem)
	fmt.Println("Client certificate and key saved to", output)
}

func writePem(filename, pemType string, bytes []byte) {
	f, _ := os.Create(filename)
	pem.Encode(f, &pem.Block{Type: pemType, Bytes: bytes})
	f.Close()
}