package proxy

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// CA — локальный корневой сертификат для перехвата HTTPS-трафика (MITM),
// по аналогии с Burp CA / mitmproxy CA. Генерируется один раз и
// сохраняется на диск; чтобы браузер перестал ругаться на HTTPS-сайты,
// ca.pem нужно один раз импортировать в доверенные корневые сертификаты
// браузера/ОС.
type CA struct {
	cert *x509.Certificate
	key  *rsa.PrivateKey

	leafCache sync.Map // host string -> *tls.Certificate
}

// CADir возвращает директорию для хранения CA (~/.burterm).
func CADir() (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", fmt.Errorf("не удалось определить домашнюю директорию: %w", err)
	}
	return filepath.Join(home, ".burterm"), nil
}

// LoadOrCreateCA загружает CA из dir/ca.pem + dir/ca-key.pem, либо создаёт
// новый и сохраняет туда же, если файлов ещё нет.
func LoadOrCreateCA(dir string) (*CA, error) {
	certPath := filepath.Join(dir, "ca.pem")
	keyPath := filepath.Join(dir, "ca-key.pem")

	certBytes, err := os.ReadFile(certPath)
	if err == nil {
		keyBytes, err := os.ReadFile(keyPath)
		if err != nil {
			return nil, fmt.Errorf("найден ca.pem, но не найден ca-key.pem рядом: %w", err)
		}
		return parseCA(certBytes, keyBytes)
	}

	return createCA(dir, certPath, keyPath)
}

func createCA(dir, certPath, keyPath string) (*CA, error) {
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return nil, fmt.Errorf("не удалось создать %s: %w", dir, err)
	}

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("генерация ключа CA: %w", err)
	}

	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("генерация серийного номера CA: %w", err)
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			CommonName:   "burterm MITM CA",
			Organization: []string{"burterm"},
		},
		NotBefore:             time.Now().Add(-24 * time.Hour),
		NotAfter:              time.Now().AddDate(10, 0, 0),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}

	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		return nil, fmt.Errorf("создание сертификата CA: %w", err)
	}

	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})

	if err := os.WriteFile(certPath, certPEM, 0o644); err != nil {
		return nil, fmt.Errorf("сохранение ca.pem: %w", err)
	}
	if err := os.WriteFile(keyPath, keyPEM, 0o600); err != nil {
		return nil, fmt.Errorf("сохранение ca-key.pem: %w", err)
	}

	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("разбор свежесозданного сертификата CA: %w", err)
	}

	return &CA{cert: cert, key: key}, nil
}

func parseCA(certPEMBytes, keyPEMBytes []byte) (*CA, error) {
	certBlock, _ := pem.Decode(certPEMBytes)
	if certBlock == nil {
		return nil, fmt.Errorf("ca.pem повреждён или пуст")
	}
	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("разбор ca.pem: %w", err)
	}

	keyBlock, _ := pem.Decode(keyPEMBytes)
	if keyBlock == nil {
		return nil, fmt.Errorf("ca-key.pem повреждён или пуст")
	}
	key, err := x509.ParsePKCS1PrivateKey(keyBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("разбор ca-key.pem: %w", err)
	}

	return &CA{cert: cert, key: key}, nil
}

// LeafFor возвращает (создавая и кэшируя при первом обращении) листовой
// TLS-сертификат для host, подписанный этим CA. Используется в
// tls.Config.GetCertificate при перехвате CONNECT-туннеля — по сертификату
// на каждый домен, с которым говорит браузер через прокси.
func (ca *CA) LeafFor(host string) (*tls.Certificate, error) {
	if cached, ok := ca.leafCache.Load(host); ok {
		return cached.(*tls.Certificate), nil
	}

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("генерация ключа для %s: %w", host, err)
	}

	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("генерация серийного номера для %s: %w", host, err)
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: host},
		DNSNames:     []string{host},
		NotBefore:    time.Now().Add(-24 * time.Hour),
		NotAfter:     time.Now().AddDate(1, 0, 0),
		KeyUsage:     x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}

	der, err := x509.CreateCertificate(rand.Reader, tmpl, ca.cert, &key.PublicKey, ca.key)
	if err != nil {
		return nil, fmt.Errorf("подпись сертификата для %s: %w", host, err)
	}

	leaf := &tls.Certificate{
		Certificate: [][]byte{der, ca.cert.Raw},
		PrivateKey:  key,
	}

	actual, _ := ca.leafCache.LoadOrStore(host, leaf)
	return actual.(*tls.Certificate), nil
}

// CertPath — путь к ca.pem, который нужно один раз импортировать
// в доверенные корневые сертификаты браузера/ОС, иначе HTTPS-сайты
// будут показывать предупреждение о недоверенном сертификате.
func (ca *CA) CertPath(dir string) string {
	return filepath.Join(dir, "ca.pem")
}