package webgateway

import (
	"crypto/tls"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/gorilla/websocket"
)

func TestConsoleUpstreamPreservesHostAndReplacesSpoofedMetadata(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Host != "console.example:9443" || r.Header.Get("X-Tera-Client-IP") != "192.0.2.77" || r.Header.Get("X-Tera-Client-Protocol") != "https" || r.Header.Get("X-Tera-Client-Port") != "9443" {
			t.Errorf("incorrect trusted client metadata: host=%s headers=%v", r.Host, r.Header)
		}
		w.Header().Set("Set-Cookie", "session=test; Secure; HttpOnly")
		_, _ = io.WriteString(w, r.URL.Path)
	}))
	defer upstream.Close()
	handler, err := NewHandler(Options{WebUpstream: upstream.URL})
	if err != nil {
		t.Fatal(err)
	}
	for _, path := range []string{"/login", "/admin/ssl", "/api/v1/config/status"} {
		r := httptest.NewRequest("GET", "https://console.example:9443"+path, nil)
		r.RemoteAddr = "192.0.2.77:5000"
		r.TLS = &tls.ConnectionState{}
		r.Header.Set("X-Tera-Client-IP", "127.0.0.1")
		r.Header.Set("X-Tera-Client-Protocol", "http")
		r.Header.Set("X-Tera-Client-Port", "80")
		w := httptest.NewRecorder()
		handler.ServeHTTP(w, r)
		if w.Code != 200 || w.Body.String() != path || !strings.Contains(w.Header().Get("Set-Cookie"), "Secure") {
			t.Fatalf("console proxy failed: %d %s", w.Code, w.Body.String())
		}
	}
}

func TestGatewayServesStaticAndRoutesRESTAndOnlyStreamWebSocketPath(t *testing.T) {
	api := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		response.Header().Set("X-Upstream", "api")
		_, _ = io.WriteString(response, request.URL.Path)
	}))
	defer api.Close()
	session := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		response.Header().Set("X-Upstream", "session")
		_, _ = io.WriteString(response, request.URL.Path)
	}))
	defer session.Close()
	staticRoot := t.TempDir()
	if err := os.WriteFile(filepath.Join(staticRoot, "index.html"), []byte("<h1>TERA</h1>"), 0o644); err != nil {
		t.Fatal(err)
	}
	handler, err := NewHandler(Options{StaticRoot: staticRoot, APIUpstream: api.URL, SessionUpstream: session.URL})
	if err != nil {
		t.Fatal(err)
	}
	server := httptest.NewServer(handler)
	defer server.Close()

	assertGatewayResponse(t, server.URL+"/api/v1/system/status", "api")
	assertGatewayResponse(t, server.URL+"/api/v1/direct/CH01/sessions", "api")
	assertGatewayResponse(t, server.URL+"/api/v1/direct/CH01/sessions/direct-1/stream", "session")
	response, err := http.Get(server.URL + "/")
	if err != nil {
		t.Fatal(err)
	}
	defer response.Body.Close()
	body, _ := io.ReadAll(response.Body)
	if !strings.Contains(string(body), "TERA") || response.Header.Get("X-Content-Type-Options") != "nosniff" {
		t.Fatalf("static response body=%q headers=%v", body, response.Header)
	}
	if policy := response.Header.Get("Content-Security-Policy"); !strings.Contains(policy, "img-src 'self' data:") {
		t.Fatalf("content security policy does not permit the in-memory raster CAPTCHA: %q", policy)
	}
}

func TestGatewayProxiesWebSocketUpgradeToSessiond(t *testing.T) {
	forwardedSources := make(chan string, 1)
	upgrader := websocket.Upgrader{
		Subprotocols: []string{"tera.direct.v1"},
		CheckOrigin:  func(*http.Request) bool { return true },
	}
	session := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		forwardedSources <- request.Header.Get("X-Forwarded-For")
		connection, err := upgrader.Upgrade(response, request, nil)
		if err != nil {
			return
		}
		defer connection.Close()
		messageType, payload, err := connection.ReadMessage()
		if err == nil {
			_ = connection.WriteMessage(messageType, payload)
		}
	}))
	defer session.Close()
	api := httptest.NewServer(http.NotFoundHandler())
	defer api.Close()
	staticRoot := t.TempDir()
	if err := os.WriteFile(filepath.Join(staticRoot, "index.html"), []byte("TERA"), 0o644); err != nil {
		t.Fatal(err)
	}
	handler, err := NewHandler(Options{StaticRoot: staticRoot, APIUpstream: api.URL, SessionUpstream: session.URL})
	if err != nil {
		t.Fatal(err)
	}
	gateway := httptest.NewServer(handler)
	defer gateway.Close()
	wsURL := "ws" + strings.TrimPrefix(gateway.URL, "http") + "/api/v1/direct/CH01/sessions/direct-1/stream"
	dialer := websocket.Dialer{Subprotocols: []string{"tera.direct.v1"}}
	connection, _, err := dialer.Dial(wsURL, http.Header{
		"Origin": []string{gateway.URL}, "X-Forwarded-For": []string{"203.0.113.200"},
	})
	if err != nil {
		t.Fatal(err)
	}
	defer connection.Close()
	want := []byte{0x00, 0x3a, 0xff}
	if err := connection.WriteMessage(websocket.BinaryMessage, want); err != nil {
		t.Fatal(err)
	}
	messageType, got, err := connection.ReadMessage()
	if err != nil || messageType != websocket.BinaryMessage || string(got) != string(want) {
		t.Fatalf("echo type=%d payload=%v err=%v", messageType, got, err)
	}
	forwarded := strings.TrimSpace(<-forwardedSources)
	if address := net.ParseIP(forwarded); address == nil || !address.IsLoopback() {
		t.Fatalf("forwarded source = %q, want actual loopback peer without spoofed chain", forwarded)
	}
}

func TestGatewayRejectsNonLoopbackUpstreams(t *testing.T) {
	if _, err := NewHandler(Options{
		StaticRoot: t.TempDir(), APIUpstream: "http://example.com", SessionUpstream: "http://127.0.0.1:8081",
	}); err == nil {
		t.Fatal("expected external API upstream rejection")
	}
}

func assertGatewayResponse(t *testing.T, url, upstream string) {
	t.Helper()
	response, err := http.Get(url)
	if err != nil {
		t.Fatal(err)
	}
	defer response.Body.Close()
	if response.Header.Get("X-Upstream") != upstream {
		t.Fatalf("%s routed to %q, want %q", url, response.Header.Get("X-Upstream"), upstream)
	}
}