package hammerhead

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	"git.veillax.com/Veillax/reef/dmarc"
	"git.veillax.com/Veillax/reef/model"
)

func reportXML(domain, id, reporter string, rows ...[4]string) string { // ip, count, dkim, spf
	var sb strings.Builder
	sb.WriteString(`<feedback><report_metadata><org_name>` + reporter + `</org_name><report_id>` + id + `</report_id><date_range><begin>` +
		strconv.FormatInt(time.Now().Add(-26*time.Hour).Unix(), 10) + `</begin><end>` + strconv.FormatInt(time.Now().Add(-2*time.Hour).Unix(), 10) + `</end></date_range></report_metadata><policy_published><domain>` + domain + `</domain><p>none</p><pct>100</pct></policy_published>`)
	for _, r := range rows {
		disp := "none"
		if r[2] != "pass" && r[3] != "pass" {
			disp = "quarantine"
		}
		sb.WriteString(`<record><row><source_ip>` + r[0] + `</source_ip><count>` + r[1] + `</count><policy_evaluated><disposition>` + disp + `</disposition><dkim>` + r[2] + `</dkim><spf>` + r[3] + `</spf></policy_evaluated></row><identifiers><header_from>` + domain + `</header_from></identifiers></record>`)
	}
	sb.WriteString(`</feedback>`)
	return sb.String()
}

func gzipped(s string) []byte {
	var b bytes.Buffer
	w := gzip.NewWriter(&b)
	w.Write([]byte(s))
	w.Close()
	return b.Bytes()
}

// dmarcFx is an org that sends as example.com from 203.0.113.7.
func dmarcFx(t *testing.T) (*fx, *orgCtx) {
	f := newFx(t)
	f.s.mu.Lock()
	f.s.st.DKIM = append(f.s.st.DKIM, &model.DKIMKey{ID: "k1", OrgID: f.org, Domain: "example.com", Selector: "s1"})
	f.s.st.MailEgress = append(f.s.st.MailEgress, &model.MailEgress{ID: "x1", OrgID: f.org, Name: "relay", SiteID: "s1", EdgeID: "e1", SourceIP: "203.0.113.7", HeloName: "mail.example.com", Enabled: true})
	org := f.s.findOrg(f.org)
	f.s.mu.Unlock()
	return f, &orgCtx{Org: *org, Actor: "test"}
}

func upload(t *testing.T, f *fx, c *orgCtx, body []byte) (int, map[string]any) {
	w := httptest.NewRecorder()
	f.s.uploadDMARC(w, httptest.NewRequest("POST", "/", bytes.NewReader(body)), c)
	var out map[string]any
	json.Unmarshal(w.Body.Bytes(), &out)
	return w.Code, out
}

func TestUploadAndSummarize(t *testing.T) {
	f, c := dmarcFx(t)
	code, out := upload(t, f, c, gzipped(reportXML("example.com", "r1", "google.com", [4]string{"203.0.113.7", "30", "pass", "pass"}, [4]string{"198.51.100.9", "5", "fail", "fail"})))
	if code != 200 || out["messages"] != float64(35) || out["duplicate"] != false {
		t.Fatalf("%d %v", code, out)
	}
	if _, out = upload(t, f, c, []byte(reportXML("example.com", "r1", "google.com"))); out["duplicate"] != true {
		t.Errorf("the same report twice should be a duplicate: %v", out)
	}
	upload(t, f, c, []byte(reportXML("example.com", "r2", "yahoo.com", [4]string{"203.0.113.7", "10", "pass", "fail"})))

	w := httptest.NewRecorder()
	f.s.listDMARC(w, httptest.NewRequest("GET", "/", nil), c)
	var doms []dmarcDomain
	if err := json.Unmarshal(w.Body.Bytes(), &doms); err != nil || len(doms) != 1 {
		t.Fatalf("%v %s", err, w.Body.String())
	}
	d := doms[0]
	if d.Reports != 2 || d.Messages != 45 || d.Passed != 40 || d.OurMsgs != 40 || d.OurPassed != 40 {
		t.Errorf("totals: %+v", d)
	}
	if len(d.Sources) != 2 || d.Sources[0].IP != "198.51.100.9" || d.Sources[0].Ours || d.Sources[0].Hint == "" {
		t.Errorf("failing sources should come first with a hint: %+v", d.Sources)
	}
	if !strings.Contains(strings.Join(d.Advice, " "), "failed") {
		t.Errorf("advice should flag the unknown sender: %v", d.Advice)
	}
}

func TestUploadRefusals(t *testing.T) {
	f, c := dmarcFx(t)
	if code, _ := upload(t, f, c, []byte("garbage")); code != 400 {
		t.Errorf("garbage: %d", code)
	}
	code, out := upload(t, f, c, []byte(reportXML("someone-else.org", "r9", "google.com", [4]string{"1.2.3.4", "1", "pass", "pass"})))
	if code != 400 || !strings.Contains(out["error"].(string), "someone-else.org") {
		t.Errorf("a domain the org doesn't send from: %d %v", code, out)
	}
	big := make([]byte, maxDMARCUpload+1)
	w := httptest.NewRecorder()
	f.s.uploadDMARC(w, httptest.NewRequest("POST", "/", bytes.NewReader(big)), c)
	if w.Code != http.StatusRequestEntityTooLarge {
		t.Errorf("oversize: %d", w.Code)
	}
}

func TestReportsStayInTheirOrg(t *testing.T) {
	f, c := dmarcFx(t)
	upload(t, f, c, []byte(reportXML("example.com", "r1", "google.com", [4]string{"203.0.113.7", "5", "pass", "pass"})))
	if n := len(f.s.db.dmarcReports("another-org", time.Time{})); n != 0 {
		t.Errorf("another org sees %d reports", n)
	}
	if n, _ := f.s.db.deleteDMARCDomain(f.org, "example.com"); n != 1 {
		t.Errorf("delete removed %d", n)
	}
}

func TestSummarizeAdvice(t *testing.T) {
	rep := func(p string, pct int, rows ...dmarc.Record) *dmarc.Report {
		return &dmarc.Report{Org: "g", ReportID: "1", Domain: "example.com", Begin: time.Unix(1, 0), End: time.Unix(2, 0), Policy: dmarc.Policy{P: p, Pct: pct}, Records: rows}
	}
	ours := map[string]bool{"203.0.113.7": true}
	pass := dmarc.Record{SourceIP: "203.0.113.7", Count: 50, DKIM: "pass", SPF: "pass"}

	got := summarizeDMARC([]*dmarc.Report{rep("none", 100, pass)}, ours)[0]
	if !strings.Contains(strings.Join(got.Advice, " "), "p=quarantine") {
		t.Errorf("a clean domain at p=none should be told to tighten: %v", got.Advice)
	}
	bad := dmarc.Record{SourceIP: "203.0.113.7", Count: 50, DKIM: "fail", SPF: "fail", Disposition: "quarantine"}
	got = summarizeDMARC([]*dmarc.Report{rep("quarantine", 50, pass, bad)}, ours)[0]
	joined := strings.Join(got.Advice, " ")
	if !strings.Contains(joined, "50%") || !strings.Contains(joined, "pct=50") {
		t.Errorf("advice: %v", got.Advice)
	}
	got = summarizeDMARC([]*dmarc.Report{rep("quarantine", 100, bad)}, ours)[0]
	if got.Sources[0].Hint == "" || !strings.Contains(got.Sources[0].Hint, "SPF and DKIM") {
		t.Errorf("own relay failing both: %q", got.Sources[0].Hint)
	}
}

func TestOwnMailFailingDMARCRaisesAnAlert(t *testing.T) {
	f, c := dmarcFx(t)
	upload(t, f, c, []byte(reportXML("example.com", "r1", "google.com", [4]string{"203.0.113.7", "40", "fail", "fail"}, [4]string{"203.0.113.7", "10", "pass", "pass"})))
	f.s.mu.Lock()
	probs := f.s.problemsLocked(f.base)
	f.s.mu.Unlock()
	p, ok := probs["dmarcrate:"+f.org+":example.com"]
	if !ok || !strings.Contains(p.msg, "20%") {
		t.Fatalf("expected a DMARC pass-rate alert, got %v", probs)
	}
	// a healthy domain raises none
	upload(t, f, c, []byte(reportXML("example.com", "r2", "yahoo.com", [4]string{"203.0.113.7", "500", "pass", "pass"})))
	f.s.mu.Lock()
	probs = f.s.problemsLocked(f.base)
	f.s.mu.Unlock()
	if _, ok := probs["dmarcrate:"+f.org+":example.com"]; ok {
		t.Errorf("alert should clear once the pass rate recovers: %v", probs)
	}
}