package security

import (
	"crypto/ed25519"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type MeshVerifier struct {
	mode   string // ed25519|hmac
	priv   ed25519.PrivateKey
	pubs   map[string]ed25519.PublicKey // nodeID->pub
	secret []byte
}

func NewMeshAuth(mode, privPath, pubDir, secret string) (*MeshVerifier, *MeshVerifier) {
	mv := &MeshVerifier{mode: mode}
	if mode == "ed25519" {
		if privPath != "" {
			if b, err := os.ReadFile(privPath); err == nil {
				if p, _ := pem.Decode(b); p != nil {
					mv.priv = ed25519.PrivateKey(p.Bytes)
				}
			}
		}
		mv.pubs = map[string]ed25519.PublicKey{}
		_ = filepath.WalkDir(pubDir, func(path string, d os.DirEntry, err error) error {
			if err != nil || d.IsDir() {
				return nil
			}
			if !strings.HasSuffix(path, ".pub") && !strings.HasSuffix(path, ".pem") {
				return nil
			}
			name := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
			b, _ := os.ReadFile(path)
			p, _ := pem.Decode(b)
			if p != nil {
				mv.pubs[name] = ed25519.PublicKey(p.Bytes)
			}
			return nil
		})
	} else {
		mv.secret = []byte(secret)
	}
	return mv, mv
}

func (m *MeshVerifier) bodyHash(r *http.Request) string {
	// Expect header X-Mesh-Hash (hex) from client for POST; for GET may be empty
	return r.Header.Get("X-Mesh-Hash")
}

func (m *MeshVerifier) canonical(r *http.Request, ts, nonce, bodyHash string) []byte {
	return []byte(strings.Join([]string{r.Method, r.URL.Path, ts, nonce, bodyHash}, "\n"))
}

// Verify incoming request (used in middleware)
func (m *MeshVerifier) VerifyIncoming(r *http.Request, now time.Time, nonces *NonceStore) bool {
	ts := r.Header.Get("X-Mesh-TS")
	nonce := r.Header.Get("X-Mesh-Nonce")
	from := r.Header.Get("X-Mesh-From")
	bh := m.bodyHash(r)
	if ts == "" || nonce == "" || from == "" {
		return false
	}
	// timestamp window
	//sec, _ := time.ParseDuration("0s")
	if len(ts) >= 10 { // crude parse
		i := int64(0)
		for _, ch := range ts {
			if ch < '0' || ch > '9' {
				i = 0
				break
			}
			i = i*10 + int64(ch-'0')
		}
		if i != 0 {
			dt := time.Unix(i, 0).Sub(now)
			if dt > 60*time.Second || dt < -60*time.Second {
				return false
			}
		}
	}
	if !nonces.Use(nonce, now) {
		return false
	}
	canon := m.canonical(r, ts, nonce, bh)
	if m.mode == "ed25519" {
		sig := r.Header.Get("X-Mesh-Sig")
		pub := m.pubs[from]
		if len(pub) == 0 {
			return false
		}
		b, err := base64.StdEncoding.DecodeString(sig)
		if err != nil {
			return false
		}
		return ed25519.Verify(pub, canon, b)
	}
	macHex := r.Header.Get("X-Mesh-MAC")
	if macHex == "" {
		return false
	}
	mac := hmac.New(sha256.New, m.secret)
	mac.Write(canon)
	want := mac.Sum(nil)
	got, _ := hex.DecodeString(macHex)
	return hmac.Equal(want, got)
}

// Sign outgoing mesh request (client-side helper)
func (m *MeshVerifier) SignOutgoing(r *http.Request, bodyHash string, from string) {
	ts := time.Now().Unix()
	nonce := make([]byte, 12)
	_, _ = rand.Read(nonce)
	nonceB64 := base64.RawURLEncoding.EncodeToString(nonce)
	r.Header.Set("X-Mesh-TS", fmtItoa(ts))
	r.Header.Set("X-Mesh-Nonce", nonceB64)
	r.Header.Set("X-Mesh-From", from)
	r.Header.Set("X-Mesh-Hash", bodyHash)
	canon := m.canonical(r, fmtItoa(ts), nonceB64, bodyHash)
	if m.mode == "ed25519" {
		sig := ed25519.Sign(m.priv, canon)
		r.Header.Set("X-Mesh-Sig", base64.StdEncoding.EncodeToString(sig))
	} else {
		mac := hmac.New(sha256.New, m.secret)
		mac.Write(canon)
		r.Header.Set("X-Mesh-MAC", hex.EncodeToString(mac.Sum(nil)))
	}
}

// Middleware enforcing mesh auth
func MeshAuthMiddleware(verifier *MeshVerifier, nonces *NonceStore, maxBody int64) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// limit headers/body
			r.Body = http.MaxBytesReader(w, r.Body, maxBody)
			if !verifier.VerifyIncoming(r, time.Now(), nonces) {
				http.Error(w, "unauthorized", http.StatusUnauthorized)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// helpers to avoid extra imports

func fmtItoa(i int64) string { // fast itoa
	if i == 0 {
		return "0"
	}
	neg := i < 0
	if neg {
		i = -i
	}
	buf := make([]byte, 0, 20)
	for i > 0 {
		buf = append(buf, byte('0'+i%10))
		i /= 10
	}
	for l, r := 0, len(buf)-1; l < r; l, r = l+1, r-1 {
		buf[l], buf[r] = buf[r], buf[l]
	}
	if neg {
		return "-" + string(buf)
	}
	return string(buf)
}

func LoadCertPool(caPath string) *x509.CertPool {
	b, err := os.ReadFile(caPath)
	if err != nil {
		return nil
	}
	pool := x509.NewCertPool()
	pool.AppendCertsFromPEM(b)
	return pool
}