package federation

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net/http"
	"strings"
	"time"

	"xorm.io/xorm"
)

// Sign adds an HTTP Signature header to req using the RSA private key.
// Implements draft-cavage-http-signatures (the fediverse de-facto standard).
// Signs: (request-target), host, date. If body is set, also signs digest.
func Sign(req *http.Request, keyID, privateKeyPEM string) error {
	if req.Header.Get("Date") == "" {
		req.Header.Set("Date", time.Now().UTC().Format(http.TimeFormat))
	}
	if req.Header.Get("Host") == "" {
		req.Header.Set("Host", req.URL.Host)
	}

	method := strings.ToLower(req.Method)
	target := req.URL.RequestURI()
	signingString := fmt.Sprintf("(request-target): %s %s\nhost: %s\ndate: %s",
		method, target,
		req.Header.Get("Host"),
		req.Header.Get("Date"),
	)
	headers := "(request-target) host date"

	priv, err := parsePrivateKey(privateKeyPEM)
	if err != nil {
		return fmt.Errorf("parse private key: %w", err)
	}

	h := sha256.Sum256([]byte(signingString))
	sig, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA256, h[:])
	if err != nil {
		return fmt.Errorf("sign: %w", err)
	}

	req.Header.Set("Signature", fmt.Sprintf(
		`keyId="%s",algorithm="rsa-sha256",headers="%s",signature="%s"`,
		keyID, headers, base64.StdEncoding.EncodeToString(sig),
	))
	return nil
}

// Verify validates the HTTP Signature header on an incoming request.
// It fetches the sender's public key from their actor document (or the local DB).
func Verify(r *http.Request, db *xorm.Engine, instanceURL string) error {
	sigHeader := r.Header.Get("Signature")
	if sigHeader == "" {
		return fmt.Errorf("missing Signature header")
	}

	params := parseSignatureHeader(sigHeader)
	keyID := params["keyId"]
	sigB64 := params["signature"]
	headersList := params["headers"]
	if keyID == "" || sigB64 == "" {
		return fmt.Errorf("malformed Signature header")
	}

	sig, err := base64.StdEncoding.DecodeString(sigB64)
	if err != nil {
		return fmt.Errorf("decode signature: %w", err)
	}

	// Fetch the public key for this keyId.
	// keyId is typically "{actorURL}#main-key" — strip the fragment to get the actor APID.
	actorAPID := strings.SplitN(keyID, "#", 2)[0]
	pubKeyPEM, err := resolvePublicKey(db, actorAPID, instanceURL)
	if err != nil {
		return fmt.Errorf("resolve public key for %s: %w", actorAPID, err)
	}

	// Reconstruct the signing string from the request.
	signedHeaders := strings.Fields(headersList)
	if len(signedHeaders) == 0 {
		signedHeaders = []string{"date"}
	}
	var parts []string
	for _, h := range signedHeaders {
		switch h {
		case "(request-target)":
			parts = append(parts, fmt.Sprintf("(request-target): %s %s",
				strings.ToLower(r.Method), r.URL.RequestURI()))
		default:
			parts = append(parts, h+": "+r.Header.Get(http.CanonicalHeaderKey(h)))
		}
	}
	signingString := strings.Join(parts, "\n")

	pub, err := parsePublicKey(pubKeyPEM)
	if err != nil {
		return fmt.Errorf("parse public key: %w", err)
	}

	h := sha256.Sum256([]byte(signingString))
	if err := rsa.VerifyPKCS1v15(pub, crypto.SHA256, h[:], sig); err != nil {
		return fmt.Errorf("signature verification failed: %w", err)
	}
	return nil
}

// ── helpers ──────────────────────────────────────────────────────────────────

func parseSignatureHeader(header string) map[string]string {
	params := make(map[string]string)
	for _, part := range strings.Split(header, ",") {
		part = strings.TrimSpace(part)
		idx := strings.Index(part, "=")
		if idx < 0 {
			continue
		}
		key := strings.TrimSpace(part[:idx])
		val := strings.Trim(strings.TrimSpace(part[idx+1:]), `"`)
		params[key] = val
	}
	return params
}

func parsePrivateKey(pemStr string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return nil, fmt.Errorf("no PEM block found")
	}
	return x509.ParsePKCS1PrivateKey(block.Bytes)
}

func parsePublicKey(pemStr string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return nil, fmt.Errorf("no PEM block found")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA public key")
	}
	return rsaPub, nil
}

// resolvePublicKey returns the public key PEM for an actor APID.
// Checks local actors first, then remote cache, then fetches from network.
func resolvePublicKey(db *xorm.Engine, actorAPID, instanceURL string) (string, error) {
	// Check if it's a local actor.
	var local struct {
		PublicKey string `xorm:"public_key"`
	}
	if found, _ := db.Table("federation_actor").
		Where("ap_id = ?", actorAPID).
		Cols("public_key").Get(&local); found && local.PublicKey != "" {
		return local.PublicKey, nil
	}

	// Fetch (and cache) from network.
	remote, err := FetchActor(db, actorAPID)
	if err != nil {
		return "", err
	}
	return remote.PublicKey, nil
}