// Package auth provides mTLS certificate helpers for IOP transport.
package auth

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"os"
	"strings"
)

// WorkloadIdentity is the authenticated role and logical enrollment name from
// one canonical IOP SPIFFE URI SAN.
type WorkloadIdentity struct {
	Role string
	Name string
}

// LoadServerTLS builds a *tls.Config for the IOP TCP server with mTLS.
// cert, key: server certificate and private key paths.
// ca: path to the CA certificate that signs client certificates.
func LoadServerTLS(cert, key, ca string) (*tls.Config, error) {
	return LoadServerTLSWithIdentity(cert, key, ca, "", "")
}

// LoadHTTPServerTLS builds a server-authenticated TLS 1.3 configuration for a
// principal-authenticated HTTPS boundary. It deliberately does not request a
// client certificate; principal authentication is performed by the handler.
func LoadHTTPServerTLS(cert, key string) (*tls.Config, error) {
	serverCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return nil, fmt.Errorf("load server cert: %w", err)
	}
	return &tls.Config{Certificates: []tls.Certificate{serverCert}, MinVersion: tls.VersionTLS13}, nil
}

// LoadServerTLSWithIdentity builds an mTLS TLS 1.3 server configuration and,
// when peerRole is non-empty, requires the client certificate to contain the
// one exact URI SAN spiffe://iop/<role>/<name>. An empty peerName accepts any
// canonical non-empty name for the required role.
func LoadServerTLSWithIdentity(cert, key, ca, peerRole, peerName string) (*tls.Config, error) {
	serverCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return nil, fmt.Errorf("load server cert: %w", err)
	}

	caCert, err := os.ReadFile(ca)
	if err != nil {
		return nil, fmt.Errorf("read CA cert: %w", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(caCert) {
		return nil, fmt.Errorf("failed to parse CA cert")
	}

	cfg := &tls.Config{
		Certificates: []tls.Certificate{serverCert},
		ClientCAs:    pool,
		ClientAuth:   tls.RequireAndVerifyClientCert,
		MinVersion:   tls.VersionTLS13,
	}
	installPeerIdentityVerifier(cfg, peerRole, peerName)
	return cfg, nil
}

// LoadClientTLS builds a *tls.Config for an IOP TCP client with mTLS.
func LoadClientTLS(cert, key, ca string) (*tls.Config, error) {
	return LoadClientTLSWithIdentity(cert, key, ca, "", "", "")
}

// LoadClientTLSWithIdentity builds an mTLS TLS 1.3 client configuration with
// normal DNS-name validation plus optional exact IOP workload identity checks.
func LoadClientTLSWithIdentity(cert, key, ca, serverName, peerRole, peerName string) (*tls.Config, error) {
	clientCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return nil, fmt.Errorf("load client cert: %w", err)
	}

	caCert, err := os.ReadFile(ca)
	if err != nil {
		return nil, fmt.Errorf("read CA cert: %w", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(caCert) {
		return nil, fmt.Errorf("failed to parse CA cert")
	}

	cfg := &tls.Config{
		Certificates: []tls.Certificate{clientCert},
		RootCAs:      pool,
		MinVersion:   tls.VersionTLS13,
		ServerName:   strings.TrimSpace(serverName),
	}
	installPeerIdentityVerifier(cfg, peerRole, peerName)
	return cfg, nil
}

func installPeerIdentityVerifier(cfg *tls.Config, role, name string) {
	role = strings.TrimSpace(role)
	name = strings.TrimSpace(name)
	if cfg == nil || role == "" {
		return
	}
	cfg.VerifyConnection = func(state tls.ConnectionState) error {
		if len(state.PeerCertificates) == 0 {
			return fmt.Errorf("peer certificate is required")
		}
		identity, err := ParseWorkloadIdentity(state.PeerCertificates[0])
		if err != nil {
			return err
		}
		if identity.Role != role || (name != "" && identity.Name != name) {
			return fmt.Errorf("peer workload identity mismatch")
		}
		return nil
	}
}

// ParseWorkloadIdentity returns the one canonical
// spiffe://iop/<role>/<name> identity in certificate. Non-IOP URI SANs are
// ignored, while missing, malformed, or multiple IOP identities are rejected.
func ParseWorkloadIdentity(certificate *x509.Certificate) (WorkloadIdentity, error) {
	if certificate == nil {
		return WorkloadIdentity{}, fmt.Errorf("peer certificate is required")
	}

	var identity *WorkloadIdentity
	for _, uri := range certificate.URIs {
		if uri == nil || uri.Scheme != "spiffe" || uri.Host != "iop" {
			continue
		}
		if identity != nil {
			return WorkloadIdentity{}, fmt.Errorf("multiple IOP workload identities")
		}
		if uri.User != nil || uri.Opaque != "" || uri.RawQuery != "" || uri.Fragment != "" || uri.ForceQuery {
			return WorkloadIdentity{}, fmt.Errorf("malformed IOP workload identity")
		}
		segments := strings.Split(strings.TrimPrefix(uri.Path, "/"), "/")
		if !strings.HasPrefix(uri.Path, "/") || len(segments) != 2 || !validWorkloadSegment(segments[0]) || !validWorkloadSegment(segments[1]) {
			return WorkloadIdentity{}, fmt.Errorf("malformed IOP workload identity")
		}
		candidate := WorkloadIdentity{Role: segments[0], Name: segments[1]}
		if uri.String() != "spiffe://iop/"+candidate.Role+"/"+candidate.Name {
			return WorkloadIdentity{}, fmt.Errorf("non-canonical IOP workload identity")
		}
		identity = &candidate
	}
	if identity == nil {
		return WorkloadIdentity{}, fmt.Errorf("IOP workload identity is required")
	}
	return *identity, nil
}

func validWorkloadSegment(value string) bool {
	return value != "" && value == strings.TrimSpace(value) && value != "." && value != ".." &&
		!strings.ContainsAny(value, "/ \t\r\n")
}

// PeerWorkloadIdentity reads the authenticated identity from a completed TLS
// connection. Enrollment handlers call it only after receiving the first
// application request, when the TLS handshake has completed.
func PeerWorkloadIdentity(conn net.Conn) (WorkloadIdentity, error) {
	tlsConn, ok := conn.(*tls.Conn)
	if !ok {
		return WorkloadIdentity{}, fmt.Errorf("TLS peer connection is required")
	}
	state := tlsConn.ConnectionState()
	if !state.HandshakeComplete {
		return WorkloadIdentity{}, fmt.Errorf("TLS peer handshake is incomplete")
	}
	if len(state.PeerCertificates) == 0 {
		return WorkloadIdentity{}, fmt.Errorf("peer certificate is required")
	}
	return ParseWorkloadIdentity(state.PeerCertificates[0])
}