// Package keymailtest keeps sign-in's Keymail classification inside the
// test process. It is imported by tests only.
//
// rastrillo/auth (at the version this app pins) builds its classifier
// with no seams: TXT lookups go through net.DefaultResolver and the
// well-known probe, the federation lookup and the token exchange go
// through http.DefaultTransport. So those two globals are the hooks,
// exactly as oficina/docs' auth_keymail_test.go uses them. Install
// replaces both for the life of the test binary:
//
//   - every DNS query is answered NXDOMAIN by a resolver on loopback,
//     so no DNS packet leaves the process, and a domain with no
//     `_keymail` record is probed directly — keymail.dev's own shape;
//   - an https request to a host registered with Serve is answered
//     in-process by that fixture; anything else falls through to the
//     real transport, where the resolver above makes a name
//     unreachable (IP literals, which httptest servers use, still work).
//
// Every address at a domain nobody registered therefore classifies as
// "not Keymail", which is what every pre-existing sign-in test assumes.
package keymailtest

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"time"
)

var (
	mu      sync.Mutex
	servers = map[string]*Server{}
)

// Install swaps net.DefaultResolver and http.DefaultTransport for the
// stubs described in the package comment and returns a function that
// puts them back. Call it once, from TestMain.
func Install() (restore func(), err error) {
	dns, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	go answerNXDOMAIN(dns)
	oldResolver, oldTransport := net.DefaultResolver, http.DefaultTransport
	net.DefaultResolver = &net.Resolver{PreferGo: true, Dial: func(ctx context.Context, _, _ string) (net.Conn, error) {
		return (&net.Dialer{}).DialContext(ctx, "udp", dns.LocalAddr().String())
	}}
	http.DefaultTransport = transport{next: oldTransport}
	return func() {
		net.DefaultResolver, http.DefaultTransport = oldResolver, oldTransport
		_ = dns.Close()
	}, nil
}

// answerNXDOMAIN echoes each query back as a "no such name" answer.
func answerNXDOMAIN(conn net.PacketConn) {
	b := make([]byte, 4096)
	for {
		n, addr, err := conn.ReadFrom(b)
		if err != nil {
			return
		}
		if n < 12 {
			continue
		}
		reply := append([]byte(nil), b[:n]...)
		reply[2], reply[3] = 0x81, 0x83 // response, recursion, NXDOMAIN
		for i := 6; i < 12; i++ {
			reply[i] = 0 // no answer, authority or additional records
		}
		_, _ = conn.WriteTo(reply, addr)
	}
}

type transport struct{ next http.RoundTripper }

func (t transport) RoundTrip(r *http.Request) (*http.Response, error) {
	mu.Lock()
	s := servers[r.URL.Host]
	mu.Unlock()
	if s == nil || r.URL.Scheme != "https" {
		return t.next.RoundTrip(r)
	}
	rec := httptest.NewRecorder()
	s.ServeHTTP(rec, r)
	res := rec.Result()
	res.Request = r
	return res, nil
}

// Server is one Keymail deployment at one host: the well-known
// document, the federation lookup for the accounts it holds, and the
// token exchange for the codes Approve hands out.
type Server struct {
	host string

	mu       sync.Mutex
	accounts map[string]bool
	codes    map[string]string
	lookups  int
	tokens   int
}

// Serve registers a Keymail deployment at host (which is also the
// address domain it answers for, since nothing publishes a `_keymail`
// record) holding the given accounts. It is unregistered when the test
// ends.
func Serve(t interface{ Cleanup(func()) }, host string, accounts ...string) *Server {
	s := &Server{host: host, accounts: map[string]bool{}, codes: map[string]string{}}
	for _, a := range accounts {
		s.accounts[strings.ToLower(a)] = true
	}
	mu.Lock()
	servers[host] = s
	mu.Unlock()
	t.Cleanup(func() {
		mu.Lock()
		delete(servers, host)
		mu.Unlock()
	})
	return s
}

// Approve is Keymail's consent screen saying yes for address: it
// returns the code the browser would carry back to /auth/callback.
func (s *Server) Approve(address string) string {
	var b [12]byte
	_, _ = rand.Read(b[:])
	code := hex.EncodeToString(b[:])
	s.mu.Lock()
	s.codes[code] = address
	s.mu.Unlock()
	return code
}

// Lookups is how many federation lookups reached this deployment.
func (s *Server) Lookups() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.lookups
}

// Tokens is how many token exchanges reached this deployment.
func (s *Server) Tokens() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.tokens
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	switch r.URL.Path {
	case "/.well-known/keymail":
		_ = json.NewEncoder(w).Encode(map[string]any{"version": "1", "reauth": true})
	case "/api/federation/lookup":
		s.mu.Lock()
		s.lookups++
		ok := s.accounts[strings.ToLower(r.URL.Query().Get("addr"))]
		s.mu.Unlock()
		if !ok {
			w.WriteHeader(http.StatusNotFound)
		}
		_ = json.NewEncoder(w).Encode(map[string]bool{"ok": ok})
	case "/api/oauth/token":
		var in map[string]string
		_ = json.NewDecoder(r.Body).Decode(&in)
		s.mu.Lock()
		s.tokens++
		address, ok := s.codes[in["code"]]
		delete(s.codes, in["code"])
		s.mu.Unlock()
		if r.Method != http.MethodPost || !ok {
			w.WriteHeader(http.StatusBadRequest)
			_ = json.NewEncoder(w).Encode(map[string]string{"error": "invalid_grant"})
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{
			"access_token": "synthetic", "token_type": "bearer",
			"address": address, "auth_time": time.Now().Unix(),
		})
	default:
		w.WriteHeader(http.StatusNotFound)
	}
}