package goproxy

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"hash/fnv"
	"math/big"
	"net"
	"os"
	"strings"
	"time"
)

// 内置默认CA证书和私钥
var (
	// 默认根证书
	defaultRootCAPem = []byte(`-----BEGIN CERTIFICATE-----
MIICJzCCAcygAwIBAgIITWWCIQf8/VIwCgYIKoZIzj0EAwIwUzEOMAwGA1UEBhMF
Q2hpbmExDzANBgNVBAgTBkZ1SmlhbjEPMA0GA1UEBxMGWGlhbWVuMRAwDgYDVQQK
EwdHb3Byb3h5MQ0wCwYDVQQDEwRNYXJzMB4XDTIyMDMyNTA1NDgwMFoXDTQyMDQy
NTA1NDgwMFowUzEOMAwGA1UEBhMFQ2hpbmExDzANBgNVBAgTBkZ1SmlhbjEPMA0G
A1UEBxMGWGlhbWVuMRAwDgYDVQQKEwdHb3Byb3h5MQ0wCwYDVQQDEwRNYXJzMFkw
EwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEf0mhVJmuTmxnLimKshdEE4+PYdxvBfQX
mRgsFV5KHHmxOrVJBFC/nDetmGowkARShWtBsX1Irm4w6i6Qk2QliKOBiTCBhjAO
BgNVHQ8BAf8EBAMCAQYwHQYDVR0lBBYwFAYIKwYBBQUHAwIGCCsGAQUFBwMBMBIG
A1UdEwEB/wQIMAYBAf8CAQIwHQYDVR0OBBYEFBI5TkWYcvUIWsBAdffs833FnBrI
MCIGA1UdEQQbMBmBF3FpbmdxaWFubHVkYW9AZ21haWwuY29tMAoGCCqGSM49BAMC
A0kAMEYCIQCk1DhW7AmIW/n/QLftQq8BHZKLevWYJ813zdrNr5kXlwIhAIVvqglY
9BkYWg4NEe/mVO4C5Vtu4FnzNU9I+rFpXVSO
-----END CERTIFICATE-----
`)
	// 默认根私钥
	defaultRootKeyPem = []byte(`-----BEGIN EC PRIVATE KEY-----
MHcCAQEEIAXeEHO0FtFqQhTvsn/DT4g3rEos97+1Nibp9RfKOKhroAoGCCqGSM49
AwEHoUQDQgAEf0mhVJmuTmxnLimKshdEE4+PYdxvBfQXmRgsFV5KHHmxOrVJBFC/
nDetmGowkARShWtBsX1Irm4w6i6Qk2QliA==
-----END EC PRIVATE KEY-----
`)
)

// 加载和初始化默认根证书
var (
	defaultRootCA  *x509.Certificate
	defaultRootKey *ecdsa.PrivateKey
)

func init() {
	var err error
	block, _ := pem.Decode(defaultRootCAPem)
	if block == nil {
		panic("解析默认根证书PEM块失败")
	}
	defaultRootCA, err = x509.ParseCertificate(block.Bytes)
	if err != nil {
		panic(fmt.Errorf("加载默认根证书失败: %s", err))
	}

	block, _ = pem.Decode(defaultRootKeyPem)
	if block == nil {
		panic("解析默认根私钥PEM块失败")
	}
	defaultRootKey, err = x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		panic(fmt.Errorf("加载默认根私钥失败: %s", err))
	}
}

// CertManager 证书管理器
type CertManager struct {
	// 证书缓存
	cache CertificateCache
	// 默认私钥，可用于多个证书共享
	defaultPrivateKey interface{} // 改为interface{}以支持不同类型的私钥
	// 默认使用ECDSA P-256曲线
	curve elliptic.Curve
	// 证书有效期（年）
	validityYears int
	// 是否使用ECDSA（否则使用RSA）
	useECDSA bool
}

// NewCertManager 创建证书管理器
func NewCertManager(cache CertificateCache, options ...CertManagerOption) *CertManager {
	manager := &CertManager{
		cache:         cache,
		curve:         elliptic.P256(), // 默认使用P-256曲线
		validityYears: 1,               // 默认证书有效期1年
		useECDSA:      true,            // 默认使用ECDSA
	}

	// 应用选项
	for _, option := range options {
		option(manager)
	}

	return manager
}

// CertManagerOption 证书管理器选项
type CertManagerOption func(*CertManager)

// WithUseECDSA 设置是否使用ECDSA（否则使用RSA）
func WithUseECDSA(useECDSA bool) CertManagerOption {
	return func(m *CertManager) {
		m.useECDSA = useECDSA
	}
}

// WithDefaultPrivateKey 设置是否使用默认私钥
func WithDefaultPrivateKey(enable bool) CertManagerOption {
	return func(m *CertManager) {
		if enable {
			if m.useECDSA {
				// 生成ECDSA私钥
				priv, err := ecdsa.GenerateKey(m.curve, rand.Reader)
				if err != nil {
					panic(fmt.Errorf("生成默认ECDSA私钥失败: %s", err))
				}
				m.defaultPrivateKey = priv
			} else {
				// 生成RSA私钥
				priv, err := rsa.GenerateKey(rand.Reader, 2048)
				if err != nil {
					panic(fmt.Errorf("生成默认RSA私钥失败: %s", err))
				}
				m.defaultPrivateKey = priv
			}
		}
	}
}

// WithCurve 设置椭圆曲线
func WithCurve(curve elliptic.Curve) CertManagerOption {
	return func(m *CertManager) {
		m.curve = curve
	}
}

// WithValidityYears 设置证书有效期（年）
func WithValidityYears(years int) CertManagerOption {
	return func(m *CertManager) {
		if years > 0 {
			m.validityYears = years
		}
	}
}

// GenerateTLSConfig 为指定主机生成TLS配置
func (m *CertManager) GenerateTLSConfig(host string) (*tls.Config, error) {
	// 处理可能的端口
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}

	// 检查证书缓存
	if m.cache != nil {
		// 检查主域名和子域名
		fields := strings.Split(host, ".")
		domains := []string{host}

		// 添加父域名
		if len(fields) > 2 {
			domains = append(domains, strings.Join(fields[1:], "."))
		}

		// 查找缓存
		for _, domain := range domains {
			if cert := m.cache.Get(domain); cert != nil {
				return &tls.Config{
					Certificates: []tls.Certificate{*cert},
				}, nil
			}
		}
	}

	// 生成新证书
	cert, err := m.GenerateCertificate(host, defaultRootCA, defaultRootKey)
	if err != nil {
		return nil, err
	}

	// 缓存证书
	if m.cache != nil {
		// 缓存主机名
		m.cache.Set(host, cert)

		// 如果是IP地址，不进行其他处理
		if net.ParseIP(host) != nil {
			return &tls.Config{
				Certificates: []tls.Certificate{*cert},
			}, nil
		}

		// 缓存域名和子域名证书
		fields := strings.Split(host, ".")
		if len(fields) >= 2 {
			// 缓存主域名
			domain := strings.Join(fields[1:], ".")
			m.cache.Set(domain, cert)
		}
	}

	return &tls.Config{
		Certificates: []tls.Certificate{*cert},
	}, nil
}

// GenerateCertificate 生成证书
func (m *CertManager) GenerateCertificate(host string, rootCA *x509.Certificate, rootKey *ecdsa.PrivateKey) (*tls.Certificate, error) {
	// 准备私钥
	var priv interface{}
	var pubKey interface{}
	var err error

	// 使用默认私钥或生成新私钥
	if m.defaultPrivateKey != nil {
		priv = m.defaultPrivateKey
	} else if m.useECDSA {
		// 生成ECDSA私钥
		ecdsaKey, err := ecdsa.GenerateKey(m.curve, rand.Reader)
		if err != nil {
			return nil, fmt.Errorf("生成ECDSA私钥失败: %s", err)
		}
		priv = ecdsaKey
		pubKey = &ecdsaKey.PublicKey
	} else {
		// 生成RSA私钥
		rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return nil, fmt.Errorf("生成RSA私钥失败: %s", err)
		}
		priv = rsaKey
		pubKey = &rsaKey.PublicKey
	}

	// 获取公钥
	if pubKey == nil {
		switch k := priv.(type) {
		case *ecdsa.PrivateKey:
			pubKey = &k.PublicKey
		case *rsa.PrivateKey:
			pubKey = &k.PublicKey
		default:
			return nil, fmt.Errorf("不支持的私钥类型")
		}
	}

	// 创建证书模板
	template := m.createCertificateTemplate(host)

	// 签名证书
	derBytes, err := x509.CreateCertificate(rand.Reader, template, rootCA, pubKey, rootKey)
	if err != nil {
		return nil, fmt.Errorf("创建证书失败: %s", err)
	}

	// 编码为PEM格式
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: derBytes,
	})

	// 编码私钥
	var keyPEM []byte
	switch k := priv.(type) {
	case *ecdsa.PrivateKey:
		privBytes, err := x509.MarshalECPrivateKey(k)
		if err != nil {
			return nil, fmt.Errorf("序列化ECDSA私钥失败: %s", err)
		}
		keyPEM = pem.EncodeToMemory(&pem.Block{
			Type:  "EC PRIVATE KEY",
			Bytes: privBytes,
		})
	case *rsa.PrivateKey:
		privBytes := x509.MarshalPKCS1PrivateKey(k)
		keyPEM = pem.EncodeToMemory(&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: privBytes,
		})
	default:
		return nil, fmt.Errorf("不支持的私钥类型")
	}

	// 创建TLS证书
	tlsCert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return nil, fmt.Errorf("创建TLS证书对失败: %s", err)
	}

	return &tlsCert, nil
}

// createCertificateTemplate 创建证书模板
func (m *CertManager) createCertificateTemplate(host string) *x509.Certificate {
	// 使用基于主机名的哈希值作为序列号
	fv := fnv.New64a()
	fv.Write([]byte(host))
	serialNumber := big.NewInt(0).SetUint64(fv.Sum64())

	// 准备模板
	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   host,
			Organization: []string{"GoProxy Dynamic CA"},
			Country:      []string{"CN"},
			Province:     []string{"GuangDong"},
			Locality:     []string{"Guangzhou"},
		},
		NotBefore:             time.Now().Add(-10 * time.Minute), // 提前10分钟生效，容忍时间偏差
		NotAfter:              time.Now().AddDate(m.validityYears, 0, 0),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		IsCA:                  false,
	}

	// 处理IP地址和域名
	ipAddr := net.ParseIP(host)
	if ipAddr != nil {
		template.IPAddresses = []net.IP{ipAddr}
	} else {
		// 移除可能的端口部分
		if strings.Contains(host, ":") {
			host = strings.Split(host, ":")[0]
		}

		// 将主机名添加到DNS名称列表
		template.DNSNames = []string{host}

		// 添加通配符域名支持
		fields := strings.Split(host, ".")
		fieldNum := len(fields)

		// 为每一级子域名添加通配符
		for i := 0; i <= (fieldNum - 2); i++ {
			wildcardDomain := "*." + strings.Join(fields[i:], ".")
			// 避免重复
			if wildcardDomain != host {
				template.DNSNames = append(template.DNSNames, wildcardDomain)
			}
		}
	}

	return template
}

// LoadCAFromFiles 从文件加载CA证书和私钥
func LoadCAFromFiles(certFile, keyFile string) (*x509.Certificate, *ecdsa.PrivateKey, error) {
	// 读取CA证书
	caCertPEM, err := os.ReadFile(certFile)
	if err != nil {
		return nil, nil, fmt.Errorf("读取CA证书文件失败: %s", err)
	}

	block, _ := pem.Decode(caCertPEM)
	if block == nil {
		return nil, nil, fmt.Errorf("解析CA证书PEM块失败")
	}

	caCert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, nil, fmt.Errorf("解析CA证书失败: %s", err)
	}

	// 读取CA私钥
	caKeyPEM, err := os.ReadFile(keyFile)
	if err != nil {
		return nil, nil, fmt.Errorf("读取CA私钥文件失败: %s", err)
	}

	block, _ = pem.Decode(caKeyPEM)
	if block == nil {
		return nil, nil, fmt.Errorf("解析CA私钥PEM块失败")
	}

	var caKey *ecdsa.PrivateKey

	// 尝试不同的私钥格式
	switch block.Type {
	case "EC PRIVATE KEY":
		caKey, err = x509.ParseECPrivateKey(block.Bytes)
		if err != nil {
			return nil, nil, fmt.Errorf("解析EC私钥失败: %s", err)
		}
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, nil, fmt.Errorf("解析PKCS8私钥失败: %s", err)
		}
		var ok bool
		caKey, ok = key.(*ecdsa.PrivateKey)
		if !ok {
			return nil, nil, fmt.Errorf("私钥不是ECDSA类型")
		}
	default:
		return nil, nil, fmt.Errorf("不支持的私钥类型: %s", block.Type)
	}

	return caCert, caKey, nil
}

// GetDefaultRootCA 获取默认根证书和私钥
func GetDefaultRootCA() (*x509.Certificate, *ecdsa.PrivateKey) {
	return defaultRootCA, defaultRootKey
}

// GenerateRootCA 生成新的根证书和私钥
func GenerateRootCA(validYears int) (*x509.Certificate, *ecdsa.PrivateKey, error) {
	// 生成私钥
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("生成根证书私钥失败: %s", err)
	}

	// 随机生成序列号
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, nil, fmt.Errorf("生成序列号失败: %s", err)
	}

	// 创建根证书模板
	notBefore := time.Now().Add(-10 * time.Minute)
	notAfter := notBefore.AddDate(validYears, 0, 0)

	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   "GoProxy Root CA",
			Organization: []string{"GoProxy"},
			Country:      []string{"CN"},
			Province:     []string{"GuangDong"},
			Locality:     []string{"Guangzhou"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLen:            2,
	}

	// 自签名
	derBytes, err := x509.CreateCertificate(rand.Reader, template, template, &priv.PublicKey, priv)
	if err != nil {
		return nil, nil, fmt.Errorf("创建根证书失败: %s", err)
	}

	// 解析生成的证书
	cert, err := x509.ParseCertificate(derBytes)
	if err != nil {
		return nil, nil, fmt.Errorf("解析生成的根证书失败: %s", err)
	}

	return cert, priv, nil
}

// SaveCertificateToFile 将证书和私钥保存到文件
func SaveCertificateToFile(cert *x509.Certificate, key *ecdsa.PrivateKey, certFile, keyFile string) error {
	// 保存证书
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert.Raw,
	})
	if err := os.WriteFile(certFile, certPEM, 0o644); err != nil {
		return fmt.Errorf("保存证书到文件失败: %s", err)
	}

	// 保存私钥
	keyBytes, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return fmt.Errorf("序列化私钥失败: %s", err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: keyBytes,
	})
	if err := os.WriteFile(keyFile, keyPEM, 0o600); err != nil {
		return fmt.Errorf("保存私钥到文件失败: %s", err)
	}

	return nil
}

// GenerateRootCA 生成根证书
func (m *CertManager) GenerateRootCA() (*x509.Certificate, interface{}, error) {
	// 生成私钥
	var priv interface{}
	var pubKey interface{}
	var err error

	if m.useECDSA {
		// 生成ECDSA私钥
		ecdsaKey, err := ecdsa.GenerateKey(m.curve, rand.Reader)
		if err != nil {
			return nil, nil, fmt.Errorf("生成ECDSA根私钥失败: %s", err)
		}
		priv = ecdsaKey
		pubKey = &ecdsaKey.PublicKey
	} else {
		// 生成RSA私钥
		rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return nil, nil, fmt.Errorf("生成RSA根私钥失败: %s", err)
		}
		priv = rsaKey
		pubKey = &rsaKey.PublicKey
	}

	// 创建根证书模板
	template := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName:   "GoProxy Root CA",
			Organization: []string{"GoProxy Authority"},
			Country:      []string{"CN"},
			Province:     []string{"GuangDong"},
			Locality:     []string{"Guangzhou"},
		},
		NotBefore:             time.Now().Add(-10 * time.Minute),
		NotAfter:              time.Now().AddDate(10, 0, 0), // 10年有效期
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLen:            1,
	}

	// 自签名
	derBytes, err := x509.CreateCertificate(rand.Reader, template, template, pubKey, priv)
	if err != nil {
		return nil, nil, fmt.Errorf("创建根证书失败: %s", err)
	}

	// 解析证书
	cert, err := x509.ParseCertificate(derBytes)
	if err != nil {
		return nil, nil, fmt.Errorf("解析生成的根证书失败: %s", err)
	}

	return cert, priv, nil
}