package auth

import (
	"crypto"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/crewjam/saml"
	"github.com/crewjam/saml/samlsp"
	"wucher/internal/config"
)

var SAMLMiddleware *samlsp.Middleware

func InitSAML(cfg config.Config) error {
	keyPair, err := tls.LoadX509KeyPair(cfg.SAMLCertFile, cfg.SAMLKeyFile)
	if err != nil {
		return fmt.Errorf("load saml keypair: %w", err)
	}
	cert, err := x509.ParseCertificate(keyPair.Certificate[0])
	if err != nil {
		return fmt.Errorf("parse saml certificate: %w", err)
	}
	signer, ok := keyPair.PrivateKey.(crypto.Signer)
	if !ok {
		return fmt.Errorf("saml private key is not a crypto.Signer")
	}

	idpMetadataBytes, err := os.ReadFile(cfg.SAMLMetadataFile)
	if err != nil {
		return fmt.Errorf("read idp metadata: %w", err)
	}

	var idpMetadata saml.EntityDescriptor
	if err := xml.Unmarshal(idpMetadataBytes, &idpMetadata); err != nil {
		return fmt.Errorf("parse idp metadata xml: %w", err)
	}

	rootURL, err := url.Parse(cfg.SAMLRootURL)
	if err != nil {
		return fmt.Errorf("invalid SAML_ROOT_URL: %w", err)
	}

	m, err := samlsp.New(samlsp.Options{
		EntityID:           cfg.SAMLRootURL,
		URL:                *rootURL,
		Key:                signer,
		Certificate:        cert,
		IDPMetadata:        &idpMetadata,
		AllowIDPInitiated:  true,
		CookieSameSite:     http.SameSiteLaxMode,
		DefaultRedirectURI: cfg.FrontendURL + "/dashboard",
	})
	if err != nil {
		return fmt.Errorf("init saml middleware: %w", err)
	}

	m.ServiceProvider.MetadataURL.Path = "/api/v1/auth/saml/metadata"
	m.ServiceProvider.AcsURL.Path = "/api/v1/auth/microsoft/callback"
	m.ServiceProvider.SloURL.Path = "/api/v1/auth/microsoft/logout"
	m.ServiceProvider.ValidateAudienceRestriction = func(assertion *saml.Assertion) error {
		if assertion == nil || assertion.Conditions == nil || len(assertion.Conditions.AudienceRestrictions) == 0 {
			return nil
		}

		expected := strings.TrimRight(firstNonEmpty(m.ServiceProvider.EntityID, m.ServiceProvider.MetadataURL.String()), "/")
		for _, audienceRestriction := range assertion.Conditions.AudienceRestrictions {
			if strings.TrimRight(audienceRestriction.Audience.Value, "/") == expected {
				return nil
			}
		}
		return fmt.Errorf("assertion Conditions AudienceRestriction does not contain %q", expected)
	}
	SAMLMiddleware = m
	return nil
}

func firstNonEmpty(values ...string) string {
	for _, v := range values {
		if strings.TrimSpace(v) != "" {
			return v
		}
	}
	return ""
}

func getIDPCertificates() ([]*x509.Certificate, error) {
	if SAMLMiddleware == nil || SAMLMiddleware.ServiceProvider.IDPMetadata == nil {
		return nil, fmt.Errorf("saml middleware not initialized")
	}

	var certificates []*x509.Certificate
	for _, idp := range SAMLMiddleware.ServiceProvider.IDPMetadata.IDPSSODescriptors {
		for _, kd := range idp.KeyDescriptors {
			for _, certData := range kd.KeyInfo.X509Data.X509Certificates {
				raw := strings.TrimSpace(certData.Data)
				if raw == "" {
					continue
				}
				der, err := base64.StdEncoding.DecodeString(raw)
				if err != nil {
					continue
				}
				cert, err := x509.ParseCertificate(der)
				if err != nil {
					continue
				}
				certificates = append(certificates, cert)
			}
		}
	}

	if len(certificates) == 0 {
		return nil, fmt.Errorf("no idp certificate in metadata")
	}
	return certificates, nil
}