// Package rules implements match-and-replace: user-defined rules that
// rewrite request/response headers as they pass through the proxy.
// Deliberately headers-only for now - see ApplyHeaders for why bodies
// are a separate, harder problem (noted as a follow-up in PLAN.md).
package rules

import (
	"bufio"
	"net/http"
	"net/textproto"
	"regexp"
	"sort"
	"strings"
)

// Rule is one match-and-replace rule.
type Rule struct {
	ID      int64
	Enabled bool
	Name    string
	Scope   string // "request" or "response"
	Part    string // "header" (only part supported so far)
	Match   string
	Replace string
	IsRegex bool
	// Position orders rule application (ascending) when several rules
	// could touch the same text.
	Position int
}

// ApplyHeaders rewrites h in place by serializing it to a raw
// "Name: value\r\n" block, running every enabled rule with Part=="header"
// over that text (in Position order), and reparsing the result. Working
// on the raw text rather than per-value substitution is what lets a rule
// add or remove a header entirely, not just rewrite an existing value -
// matching how Burp's header match/replace works. If a rule's output
// doesn't parse back as valid headers, ApplyHeaders returns h unchanged
// rather than risk sending something corrupted.
func ApplyHeaders(h http.Header, rs []Rule) http.Header {
	keys := make([]string, 0, len(h))
	for k := range h {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var block strings.Builder
	for _, k := range keys {
		for _, v := range h[k] {
			block.WriteString(k)
			block.WriteString(": ")
			block.WriteString(v)
			block.WriteString("\r\n")
		}
	}
	text := block.String()

	changed := false
	for _, r := range sortedByPosition(rs) {
		if !r.Enabled || r.Part != "header" {
			continue
		}
		if next, ok := apply(text, r); ok {
			text, changed = next, true
		}
	}
	if !changed {
		return h
	}

	tp := textproto.NewReader(bufio.NewReader(strings.NewReader(text + "\r\n")))
	mh, err := tp.ReadMIMEHeader()
	if err != nil {
		return h
	}
	return http.Header(mh)
}

func sortedByPosition(rs []Rule) []Rule {
	out := make([]Rule, len(rs))
	copy(out, rs)
	sort.SliceStable(out, func(i, j int) bool { return out[i].Position < out[j].Position })
	return out
}

func apply(text string, r Rule) (string, bool) {
	if r.IsRegex {
		re, err := regexp.Compile(r.Match)
		if err != nil {
			return text, false
		}
		return re.ReplaceAllString(text, r.Replace), true
	}
	if r.Match == "" {
		return text, false
	}
	return strings.ReplaceAll(text, r.Match, r.Replace), true
}