package proxy

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"
	"time"

	"git.t5de.dev/t5de/t5de-proxy/internal/imq"
)

const token = "secret"

var auth = "Basic " + base64.StdEncoding.EncodeToString([]byte("t5de:"+token))

// send makes a raw request to the proxy, as the client does, and returns the response.
func send(t *testing.T, proxy *httptest.Server, raw string) *http.Response {
	t.Helper()
	conn, err := net.Dial("tcp", proxy.Listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { conn.Close() })
	if _, err := io.WriteString(conn, raw); err != nil {
		t.Fatal(err)
	}
	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { resp.Body.Close() })
	return resp
}

func body(t *testing.T, resp *http.Response) string {
	t.Helper()
	b, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	return string(b)
}

// logBuffer collects the proxy's log, written from its handlers' goroutines.
type logBuffer struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

func (b *logBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.Write(p)
}

func (b *logBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.String()
}

func newProxy(t *testing.T, s *Server) (*httptest.Server, *logBuffer) {
	logs := &logBuffer{}
	s.Token = token
	s.Log = log.New(logs, "", 0)
	p := httptest.NewServer(s)
	t.Cleanup(p.Close)
	return p, logs
}

func TestRefusesRequestsWithoutTheToken(t *testing.T) {
	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("an unauthorized request reached the server")
	}))
	defer origin.Close()
	p, _ := newProxy(t, &Server{})
	for _, header := range []string{"", "Proxy-Authorization: Basic " + base64.StdEncoding.EncodeToString([]byte("t5de:wrong")) + "\r\n"} {
		resp := send(t, p, "GET "+origin.URL+"/ HTTP/1.1\r\nHost: x\r\n"+header+"\r\n")
		if resp.StatusCode != http.StatusProxyAuthRequired {
			t.Fatalf("got %d, want 407", resp.StatusCode)
		}
	}
}

func TestForwardsAsTheClientAsked(t *testing.T) {
	var got *http.Request
	var gotBody string
	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r
		b, _ := io.ReadAll(r.Body)
		gotBody = string(b)
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Set-Cookie", "a=b")
		w.Header().Set("Location", "https://example.invalid/next")
		w.WriteHeader(http.StatusFound)
		io.WriteString(w, "compressed bytes")
	}))
	defer origin.Close()
	p, logs := newProxy(t, &Server{})
	resp := send(t, p, "POST "+origin.URL+"/api/xmlrpc/client.php?sid=1 HTTP/1.1\r\nHost: ignored\r\n"+
		"Proxy-Authorization: "+auth+"\r\nX-imvu-auth: abc\r\nAccept-Encoding: gzip\r\nConnection: close\r\n"+
		"Content-Length: 5\r\n\r\nhello")
	// The response comes back as the server sent it: the redirect not followed, the body not decompressed.
	if resp.StatusCode != http.StatusFound || resp.Header.Get("Location") != "https://example.invalid/next" ||
		resp.Header.Get("Content-Encoding") != "gzip" || resp.Header.Get("Set-Cookie") != "a=b" {
		t.Fatalf("response: %d %v", resp.StatusCode, resp.Header)
	}
	if b := body(t, resp); b != "compressed bytes" {
		t.Fatalf("body %q", b)
	}
	if got.URL.RequestURI() != "/api/xmlrpc/client.php?sid=1" || gotBody != "hello" || got.Header.Get("X-imvu-auth") != "abc" ||
		got.Header.Get("Accept-Encoding") != "gzip" {
		t.Fatalf("request: %s %v %q", got.URL, got.Header, gotBody)
	}
	// Neither the proxy's token nor a user agent the client didn't send reach the server.
	if got.Header.Get("Proxy-Authorization") != "" || got.Header.Get("User-Agent") != "" {
		t.Fatalf("request headers: %v", got.Header)
	}
	if got.Host != strings.TrimPrefix(origin.URL, "http://") {
		t.Fatalf("Host %q", got.Host)
	}
	// The log leaves out queries.
	if !strings.Contains(logs.String(), "/api/xmlrpc/client.php?… 302") || strings.Contains(logs.String(), "sid=1") {
		t.Fatalf("log: %s", logs)
	}
}

func TestForwardsHTTPS(t *testing.T) {
	origin := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "secure "+r.URL.Path)
	}))
	defer origin.Close()
	p, _ := newProxy(t, &Server{Transport: origin.Client().Transport})
	resp := send(t, p, "GET "+origin.URL+"/x HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\n\r\n")
	if b := body(t, resp); resp.StatusCode != 200 || b != "secure /x" {
		t.Fatalf("%d %q", resp.StatusCode, b)
	}
}

// recorder is a transport that records requests and answers them itself.
type recorder struct {
	mu   sync.Mutex
	urls []string
	fail func(*http.Request) error
}

func (r *recorder) RoundTrip(req *http.Request) (*http.Response, error) {
	r.mu.Lock()
	r.urls = append(r.urls, req.URL.String())
	r.mu.Unlock()
	if r.fail != nil {
		if err := r.fail(req); err != nil {
			return nil, err
		}
	}
	return &http.Response{StatusCode: 200, Header: http.Header{}, Body: io.NopCloser(strings.NewReader("ok")), Request: req}, nil
}

func TestUpgradesIMVUHostsToHTTPS(t *testing.T) {
	rec := &recorder{}
	p, _ := newProxy(t, &Server{Transport: rec})
	for _, u := range []string{"http://client-dynamic.imvu.com/api/xmlrpc/client.php", "http://userimages03-akm.imvu.com/a.png", "http://example.com/a", "http://www.imvu.com:8080/a", "http://product.virtual.imvu.com/a"} {
		body(t, send(t, p, "GET "+u+" HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\n\r\n"))
	}
	want := []string{"https://client-dynamic.imvu.com/api/xmlrpc/client.php", "https://userimages03-akm.imvu.com/a.png", "http://example.com/a", "http://www.imvu.com:8080/a", "http://product.virtual.imvu.com/a"}
	if fmt.Sprint(rec.urls) != fmt.Sprint(want) {
		t.Fatalf("got %v, want %v", rec.urls, want)
	}
}

func TestFallsBackToHTTPWhenHTTPSFails(t *testing.T) {
	rec := &recorder{fail: func(r *http.Request) error {
		if r.URL.Scheme == "https" {
			return errors.New("tls: handshake failure")
		}
		return nil
	}}
	p, _ := newProxy(t, &Server{Transport: rec})
	resp := send(t, p, "GET http://userimages.imvu.com/a.png HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\n\r\n")
	if b := body(t, resp); resp.StatusCode != 200 || b != "ok" {
		t.Fatalf("%d %q", resp.StatusCode, b)
	}
	want := []string{"https://userimages.imvu.com/a.png", "http://userimages.imvu.com/a.png"}
	if fmt.Sprint(rec.urls) != fmt.Sprint(want) {
		t.Fatalf("got %v, want %v", rec.urls, want)
	}
}

func TestBlocksReportingEndpoints(t *testing.T) {
	rec := &recorder{}
	p, logs := newProxy(t, &Server{Transport: rec})
	for _, u := range []string{"http://bugs.imvu.com/x", "https://www.imvu.com/catalog/web_submit_logs.php", "https://api.imvu.com/events/1"} {
		if resp := send(t, p, "POST "+u+" HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\nContent-Length: 0\r\n\r\n"); resp.StatusCode != http.StatusForbidden {
			t.Fatalf("%s: %d", u, resp.StatusCode)
		}
	}
	if len(rec.urls) != 0 || strings.Count(logs.String(), "blocked") != 3 {
		t.Fatalf("forwarded %v; log %s", rec.urls, logs)
	}
	for _, u := range []string{"https://api.imvu.com/eventsx", "https://api.imvu.com/user/1"} {
		parsed, _ := url.Parse(u)
		if Blocked(parsed) {
			t.Fatalf("%s blocked", u)
		}
	}
}

func TestRejectsOriginFormRequests(t *testing.T) {
	p, _ := newProxy(t, &Server{Transport: &recorder{}})
	if resp := send(t, p, "GET /a HTTP/1.1\r\nHost: www.imvu.com\r\nProxy-Authorization: "+auth+"\r\n\r\n"); resp.StatusCode != http.StatusBadRequest {
		t.Fatalf("got %d", resp.StatusCode)
	}
}

// imqFrame encodes an IMQ frame.
func imqFrame(typ uint64, data string) []byte {
	b := []byte{0x08, byte(typ) | 0x80, byte(typ >> 7), 0x12, byte(len(data))}
	return append(b, data...)
}

func TestConnectSpeaksTLSUpstream(t *testing.T) {
	// A TLS echo server stands in for IMQ's gateway.
	gateway := httptest.NewUnstartedServer(nil)
	gateway.StartTLS()
	gateway.Close() // only for its certificate
	cert := gateway.TLS.Certificates[0]
	ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{Certificates: []tls.Certificate{cert}})
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	go func() {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		defer c.Close()
		io.Copy(c, c)
	}()

	roots := gateway.Client().Transport.(*http.Transport).TLSClientConfig.RootCAs
	p, logs := newProxy(t, &Server{TLSConfig: &tls.Config{RootCAs: roots}, Trace: true,
		Dial: func(ctx context.Context, network, addr string) (net.Conn, error) {
			if addr != "example.com:443" {
				return nil, fmt.Errorf("dialed %s", addr)
			}
			return net.Dial(network, ln.Addr().String())
		}})
	conn, err := net.Dial("tcp", p.Listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	// The test certificate is for example.com. The client's first frame comes with its CONNECT.
	hello := imqFrame(1302, "connect")
	fmt.Fprintf(conn, "CONNECT example.com:443 HTTP/1.1\r\nProxy-authorization: %s\r\n\r\n%s", auth, hello)
	r := bufio.NewReader(conn)
	status, err := r.ReadString('\n')
	if err != nil {
		t.Fatalf("%v; log %s", err, logs)
	}
	if !strings.HasPrefix(status, "HTTP/1.1 200") {
		t.Fatalf("status %q; log %s", status, logs)
	}
	for line, _ := r.ReadString('\n'); line != "\r\n"; line, _ = r.ReadString('\n') {
	}
	echo := make([]byte, len(hello))
	if _, err := io.ReadFull(r, echo); err != nil || !bytes.Equal(echo, hello) {
		t.Fatalf("echo %q %v", echo, err)
	}
	conn.Close()
	// The log names the frames each way.
	for i := 0; i < 200 && !strings.Contains(logs.String(), "closed"); i++ {
		time.Sleep(10 * time.Millisecond)
	}
	if !strings.Contains(logs.String(), "c2g C2gConnect 7B") || !strings.Contains(logs.String(), "g2c C2gConnect 7B") {
		t.Fatalf("log: %s", logs)
	}
}

func TestReportsUnreachableServers(t *testing.T) {
	rec := &recorder{fail: func(*http.Request) error { return errors.New("dial tcp: no such host") }}
	p, _ := newProxy(t, &Server{Transport: rec})
	resp := send(t, p, "GET https://nowhere.imvu.com/ HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\n\r\n")
	if resp.StatusCode != http.StatusBadGateway || !strings.Contains(resp.Header.Get(ErrorHeader), "no such host") {
		t.Fatalf("%d %v", resp.StatusCode, resp.Header)
	}
}

func TestBuffersBodiesToSendAgain(t *testing.T) {
	// Go's transport sends a request again when the server closed a pooled connection as it went out, if it can
	// get the body again (GetBody); without it, it fails with "cannot rewind body after connection loss".
	in := httptest.NewRequest("POST", "http://x/", strings.NewReader("hello"))
	out := in.Clone(in.Context())
	if ok, err := bufferBody(out, in); !ok || err != nil || out.GetBody == nil || out.ContentLength != 5 {
		t.Fatalf("%v %v %d", ok, err, out.ContentLength)
	}
	for i := 0; i < 2; i++ {
		body, _ := out.GetBody()
		if b, _ := io.ReadAll(body); string(b) != "hello" {
			t.Fatalf("body %q", b)
		}
	}
	// Bodies over maxBuffered are streamed whole, and can't be sent again.
	big := strings.Repeat("x", maxBuffered+10)
	in = httptest.NewRequest("POST", "http://x/", strings.NewReader(big))
	out = in.Clone(in.Context())
	if ok, err := bufferBody(out, in); ok || err != nil || out.GetBody != nil {
		t.Fatalf("%v %v", ok, err)
	}
	if b, _ := io.ReadAll(out.Body); len(b) != len(big) {
		t.Fatalf("streamed %d bytes", len(b))
	}
}

func TestFallsBackToHTTPWithTheBody(t *testing.T) {
	var got []string
	rec := &recorder{fail: func(r *http.Request) error {
		b, _ := io.ReadAll(r.Body)
		got = append(got, r.URL.Scheme+" "+string(b))
		if r.URL.Scheme == "https" {
			return errors.New("tls: handshake failure")
		}
		return nil
	}}
	p, _ := newProxy(t, &Server{Transport: rec})
	resp := send(t, p, "POST http://client-dynamic.imvu.com/api/xmlrpc/client.php HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\nContent-Length: 4\r\n\r\nbody")
	if resp.StatusCode != 200 || fmt.Sprint(got) != "[https body http body]" {
		t.Fatalf("%d %q", resp.StatusCode, got)
	}
}

// suffixFilter appends "!" to messages on mount "messages" and drops those saying "drop".
type suffixFilter struct{}

func (suffixFilter) Outgoing(queue, mount string, message []byte) ([]byte, bool) {
	if string(message) == "drop" {
		return nil, true
	}
	if mount == "messages" {
		return append(bytes.Clone(message), '!'), false
	}
	return message, false
}

func TestConnectFiltersOutgoingMessages(t *testing.T) {
	gateway := httptest.NewUnstartedServer(nil)
	gateway.StartTLS()
	gateway.Close()
	ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{Certificates: gateway.TLS.Certificates})
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	received := make(chan []byte, 1)
	go func() {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		defer c.Close()
		var got []byte
		buf := make([]byte, 4096)
		c.SetReadDeadline(time.Now().Add(2 * time.Second))
		for {
			n, err := c.Read(buf)
			got = append(got, buf[:n]...)
			if err != nil {
				break
			}
		}
		received <- got
	}()
	roots := gateway.Client().Transport.(*http.Transport).TLSClientConfig.RootCAs
	p, logs := newProxy(t, &Server{TLSConfig: &tls.Config{RootCAs: roots}, IMQFilter: suffixFilter{},
		Dial: func(ctx context.Context, network, addr string) (net.Conn, error) {
			return net.Dial(network, ln.Addr().String())
		}})
	conn, err := net.Dial("tcp", p.Listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	send := func(op uint64, mount, msg string) []byte {
		return imq.EncodeFrame(imq.C2gSendMessage, imq.SendMessage{OpID: op, Queue: "/chat/1", Mount: mount, Message: []byte(msg)}.Marshal())
	}
	// Sent in pieces, to check frames are put back together.
	stream := append(append(send(1, "messages", "hi"), send(2, "messages", "drop")...), send(3, "control", "x")...)
	fmt.Fprintf(conn, "CONNECT example.com:443 HTTP/1.1\r\nProxy-authorization: %s\r\n\r\n", auth)
	r := bufio.NewReader(conn)
	for line, _ := r.ReadString('\n'); line != "\r\n"; line, _ = r.ReadString('\n') {
		if line == "" {
			t.Fatalf("no response; log %s", logs)
		}
	}
	for i := 0; i < len(stream); i += 5 {
		conn.Write(stream[i:min(i+5, len(stream))])
		time.Sleep(time.Millisecond)
	}
	// The dropped message is answered by the proxy.
	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	var fromProxy []byte
	buf := make([]byte, 64)
	for len(fromProxy) < len(imq.EncodeFrame(imq.G2cResult, imq.Result(2, 0))) {
		n, err := r.Read(buf)
		if err != nil {
			t.Fatalf("%v; log %s", err, logs)
		}
		fromProxy = append(fromProxy, buf[:n]...)
	}
	if !bytes.Equal(fromProxy, imq.EncodeFrame(imq.G2cResult, imq.Result(2, 0))) {
		t.Fatalf("client got %x", fromProxy)
	}
	conn.Close()
	got := <-received
	want := append(send(1, "messages", "hi!"), send(3, "control", "x")...)
	if !bytes.Equal(got, want) {
		t.Fatalf("gateway got %q, want %q", got, want)
	}
}

// testInterceptor answers /local itself, uppercases request bodies to /change, and appends to their responses.
type testInterceptor struct{}

func (testInterceptor) Intercept(x *Exchange) error {
	switch x.URL.Path {
	case "/local":
		x.Response = LocalResponse(200, "text/plain", []byte("from the proxy"))
	case "/change":
		x.Body = bytes.ToUpper(x.Body)
		x.ModifyResponse = func(resp *http.Response, body []byte) ([]byte, error) {
			return append(body, " (changed)"...), nil
		}
	}
	return nil
}

func TestInterceptor(t *testing.T) {
	var gotBody string
	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/local" {
			t.Error("a request answered by the proxy reached the server")
		}
		b, _ := io.ReadAll(r.Body)
		gotBody = string(b)
		// Compressed, as IMVU's servers answer: the interceptor sees it decompressed.
		w.Header().Set("Content-Encoding", "gzip")
		zw := gzip.NewWriter(w)
		io.WriteString(zw, "echo "+gotBody)
		zw.Close()
	}))
	defer origin.Close()
	p, _ := newProxy(t, &Server{Interceptor: testInterceptor{}})
	resp := send(t, p, "GET "+origin.URL+"/local HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\n\r\n")
	if b := body(t, resp); b != "from the proxy" {
		t.Fatalf("local: %q", b)
	}
	resp = send(t, p, "POST "+origin.URL+"/change HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\nContent-Length: 5\r\n\r\nhello")
	if b := body(t, resp); b != "echo HELLO (changed)" || gotBody != "HELLO" || resp.Header.Get("Content-Encoding") != "" {
		t.Fatalf("changed: %q, server got %q, %v", b, gotBody, resp.Header)
	}
}