package waf

import (
	"context"
	"log/slog"
	"net/http"
	"strings"

	"github.com/dropmorepackets/haproxy-go/pkg/encoding"
	"github.com/dropmorepackets/haproxy-go/spop"
)

// SPOEAgent wraps the haproxy-go SPOE server and dispatches each
// inspected request to the appropriate per-domain Coraza engine.
type SPOEAgent struct {
	Manager *Manager
	Addr    string
}

// ListenAndServe starts the SPOE agent. Blocks until ctx is cancelled.
func (a *SPOEAgent) ListenAndServe(ctx context.Context) error {
	agent := spop.Agent{
		Addr:        a.Addr,
		Handler:     spop.HandlerFunc(a.handle),
		BaseContext: ctx,
	}
	return agent.ListenAndServe()
}

// handle is called by the haproxy-go SPOE library for every NOTIFY
// frame HAProxy sends. It extracts the request data, runs Coraza,
// and optionally sets a txn.waf.status variable to trigger a deny ACL.
func (a *SPOEAgent) handle(ctx context.Context, w *encoding.ActionWriter, m *encoding.Message) {
	var (
		clientIP string
		method   string
		uri      string // full request URI (path + optional ?query)
		httpVer  string
		host     string
		rawHdrs  string
	)

	// Iterate over the key-value pairs HAProxy sent with this message.
	entry := encoding.AcquireKVEntry()
	defer encoding.ReleaseKVEntry(entry)
	for m.KV.Next(entry) {
		switch {
		case entry.NameEquals("src"):
			addr := entry.ValueAddr()
			if addr.IsValid() {
				clientIP = addr.String()
			}
		case entry.NameEquals("method"):
			method = string(entry.ValueBytes())
		case entry.NameEquals("uri"):
			uri = string(entry.ValueBytes())
		case entry.NameEquals("ver"):
			httpVer = string(entry.ValueBytes())
		case entry.NameEquals("host"):
			host = string(entry.ValueBytes())
		case entry.NameEquals("headers"):
			rawHdrs = string(entry.ValueBytes())
		}
		entry.Reset()
	}

	if host == "" {
		return
	}

	de, ok := a.Manager.GetForHost(host)
	if !ok {
		return // WAF not configured or disabled for this domain
	}

	tx := de.WAF.NewTransaction()
	defer func() {
		tx.ProcessLogging()
		if err := tx.Close(); err != nil {
			slog.Warn("waf: tx.Close", "error", err)
		}
	}()

	// Feed connection metadata.
	if clientIP != "" {
		tx.ProcessConnection(clientIP, 0, "", 0)
	}

	if uri == "" {
		uri = "/"
	}
	if httpVer == "" {
		httpVer = "HTTP/1.1"
	}
	tx.ProcessURI(uri, method, httpVer)

	// Feed Host header first (required by many CRS rules).
	tx.AddRequestHeader("Host", host)

	// Parse and feed all raw headers.
	parseHeaders(rawHdrs, func(name, val string) {
		if !strings.EqualFold(name, "host") { // already added above
			tx.AddRequestHeader(name, val)
		}
	})

	// Evaluate request headers.
	interruption := tx.ProcessRequestHeaders()
	if interruption != nil {
		status := interruption.Status
		if status == 0 {
			status = http.StatusForbidden
		}
		slog.Info("waf: request blocked",
			"host", host,
			"method", method,
			"uri", uri,
			"client", clientIP,
			"status", status,
			"mode", de.Mode,
		)
		if de.Mode == "blocking" {
			if err := w.SetInt64(encoding.VarScopeTransaction, "status", int64(status)); err != nil {
				slog.Warn("waf: SetInt64 status", "error", err)
			}
		}
		return
	}

	// Alert-only log for detection mode.
	if tx.IsInterrupted() && de.Mode != "blocking" {
		slog.Info("waf: request flagged (detection)",
			"host", host, "method", method, "uri", uri, "client", clientIP,
		)
	}
}

// parseHeaders splits HAProxy raw headers ("Name: value\r\n…") and
// calls fn for each valid header line.
func parseHeaders(raw string, fn func(name, val string)) {
	for _, line := range strings.Split(raw, "\n") {
		line = strings.TrimRight(line, "\r")
		if line == "" {
			continue
		}
		idx := strings.IndexByte(line, ':')
		if idx <= 0 {
			continue
		}
		name := strings.TrimSpace(line[:idx])
		val := strings.TrimSpace(line[idx+1:])
		if name != "" {
			fn(name, val)
		}
	}
}