package proxy

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/kubestellar/hive/v2/pkg/agent"
)

// ---------- extractSNI: various truncation boundary tests ----------

func TestExtractSNISessionIDTruncatesCipherSuites(t *testing.T) {
	// Session ID so large that there's no room for cipher suites length
	chBody := make([]byte, 0, 40)
	chBody = append(chBody, 0x03, 0x03) // version
	chBody = append(chBody, make([]byte, 32)...) // random
	chBody = append(chBody, 2)  // session ID length = 2
	chBody = append(chBody, 0xAA) // only 1 byte of session ID (truncated)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("truncated after session ID should return empty, got %q", sni)
	}
}

func TestExtractSNITruncatedCipherSuites(t *testing.T) {
	// Build a ClientHello where cipher suites length points past the data
	chBody := make([]byte, 0, 128)
	chBody = append(chBody, 0x03, 0x03) // version
	chBody = append(chBody, make([]byte, 32)...) // random
	chBody = append(chBody, 0) // session ID len = 0
	chBody = append(chBody, 0x00, 0xFF) // cipher suites length = 255 (way past data)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("truncated cipher suites should return empty, got %q", sni)
	}
}

func TestExtractSNITruncatedCompression(t *testing.T) {
	// Valid through cipher suites, but compression length goes past data
	chBody := make([]byte, 0, 128)
	chBody = append(chBody, 0x03, 0x03)
	chBody = append(chBody, make([]byte, 32)...)
	chBody = append(chBody, 0) // session ID
	chBody = append(chBody, 0x00, 0x02, 0x00, 0x2f) // cipher suites (2 bytes)
	chBody = append(chBody, 0xFF) // compression methods length = 255 (past data)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("truncated compression should return empty, got %q", sni)
	}
}

func TestExtractSNITruncatedExtensionsLen(t *testing.T) {
	// Valid through compression, but extensions length field is truncated
	chBody := make([]byte, 0, 128)
	chBody = append(chBody, 0x03, 0x03)
	chBody = append(chBody, make([]byte, 32)...)
	chBody = append(chBody, 0) // session ID
	chBody = append(chBody, 0x00, 0x02, 0x00, 0x2f) // cipher suites
	chBody = append(chBody, 0x01, 0x00) // compression
	chBody = append(chBody, 0x00) // only 1 byte of extensions length (need 2)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("truncated extensions length should return empty, got %q", sni)
	}
}

func TestExtractSNITruncatedExtData(t *testing.T) {
	// Extensions length says more than available
	chBody := make([]byte, 0, 128)
	chBody = append(chBody, 0x03, 0x03)
	chBody = append(chBody, make([]byte, 32)...)
	chBody = append(chBody, 0)
	chBody = append(chBody, 0x00, 0x02, 0x00, 0x2f)
	chBody = append(chBody, 0x01, 0x00)
	// Extensions: say 100 bytes but only provide 4
	chBody = append(chBody, 0x00, 0x64) // extensions length = 100
	chBody = append(chBody, 0x00, 0x01) // extension type (not SNI)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("truncated extension data should return empty, got %q", sni)
	}
}

func TestExtractSNIShortSNIPayload(t *testing.T) {
	// SNI extension with payload too short (< 2 bytes for list length)
	chBody := make([]byte, 0, 128)
	chBody = append(chBody, 0x03, 0x03)
	chBody = append(chBody, make([]byte, 32)...)
	chBody = append(chBody, 0)
	chBody = append(chBody, 0x00, 0x02, 0x00, 0x2f)
	chBody = append(chBody, 0x01, 0x00)

	// SNI extension with 1-byte payload
	ext := []byte{0x00, 0x00, 0x00, 0x01, 0xFF} // type=SNI, len=1, data=0xFF
	chBody = append(chBody, byte(len(ext)>>8), byte(len(ext)))
	chBody = append(chBody, ext...)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("short SNI payload should return empty, got %q", sni)
	}
}

func TestExtractSNINonHostNameType(t *testing.T) {
	// SNI extension with a non-host_name type (nameType != 0)
	chBody := make([]byte, 0, 200)
	chBody = append(chBody, 0x03, 0x03)
	chBody = append(chBody, make([]byte, 32)...)
	chBody = append(chBody, 0)
	chBody = append(chBody, 0x00, 0x02, 0x00, 0x2f)
	chBody = append(chBody, 0x01, 0x00)

	// SNI extension with nameType=1 (not host_name=0)
	hostname := []byte("example.com")
	sniListLen := 1 + 2 + len(hostname)
	sniPayload := make([]byte, 0, 20+len(hostname))
	sniPayload = append(sniPayload, byte(sniListLen>>8), byte(sniListLen))
	sniPayload = append(sniPayload, 0x01) // nameType=1 (NOT host_name)
	sniPayload = append(sniPayload, byte(len(hostname)>>8), byte(len(hostname)))
	sniPayload = append(sniPayload, hostname...)

	ext := make([]byte, 0, 4+len(sniPayload))
	ext = append(ext, 0x00, 0x00) // SNI type
	ext = append(ext, byte(len(sniPayload)>>8), byte(len(sniPayload)))
	ext = append(ext, sniPayload...)

	chBody = append(chBody, byte(len(ext)>>8), byte(len(ext)))
	chBody = append(chBody, ext...)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("non-host_name type should return empty, got %q", sni)
	}
}

func TestExtractSNITruncatedName(t *testing.T) {
	// SNI extension where nameLen extends past the data
	chBody := make([]byte, 0, 128)
	chBody = append(chBody, 0x03, 0x03)
	chBody = append(chBody, make([]byte, 32)...)
	chBody = append(chBody, 0)
	chBody = append(chBody, 0x00, 0x02, 0x00, 0x2f)
	chBody = append(chBody, 0x01, 0x00)

	// SNI extension: nameLen says 100 but only 3 bytes of name available
	sniPayload := []byte{
		0x00, 0x06, // list len = 6
		0x00,       // nameType = host_name
		0x00, 0x64, // nameLen = 100 (way past available)
		'a', 'b', 'c',
	}
	ext := make([]byte, 0, 4+len(sniPayload))
	ext = append(ext, 0x00, 0x00) // SNI type
	ext = append(ext, byte(len(sniPayload)>>8), byte(len(sniPayload)))
	ext = append(ext, sniPayload...)

	chBody = append(chBody, byte(len(ext)>>8), byte(len(ext)))
	chBody = append(chBody, ext...)

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("truncated name should return empty, got %q", sni)
	}
}

func TestExtractSNITruncatedSessionID(t *testing.T) {
	// After random (34 bytes), session ID length byte goes past data
	chBody := make([]byte, 0, 35)
	chBody = append(chBody, 0x03, 0x03)     // version
	chBody = append(chBody, make([]byte, 32)...) // random
	// No session ID length byte — end of data

	hsHeader := []byte{0x01}
	hsLen := len(chBody)
	hsHeader = append(hsHeader, byte(hsLen>>16), byte(hsLen>>8), byte(hsLen))
	handshake := append(hsHeader, chBody...)

	record := []byte{0x16, 0x03, 0x01}
	record = append(record, byte(len(handshake)>>8), byte(len(handshake)))
	record = append(record, handshake...)

	sni := extractSNI(record)
	if sni != "" {
		t.Errorf("no session ID should return empty, got %q", sni)
	}
}

// ---------- handleConnectDirect: SplitHostPort error ----------

func TestHandleConnectDirectNoPort(t *testing.T) {
	p := newTestProxy()
	p.inference = newInferenceRouter()

	clientConn, proxyClient := net.Pipe()
	defer clientConn.Close()

	req, _ := http.NewRequest("CONNECT", "http://api.github.com", nil)
	req.Host = "api.github.com" // no port — SplitHostPort will fail, host remains api.github.com

	go func() {
		p.handleConnectDirect(proxyClient, req)
		proxyClient.Close()
	}()

	// Read the "200 Connection established"
	buf := make([]byte, 4096)
	var collected []byte
	for {
		clientConn.SetReadDeadline(time.Now().Add(3 * time.Second))
		n, err := clientConn.Read(buf)
		if n > 0 {
			collected = append(collected, buf[:n]...)
		}
		if err != nil {
			break
		}
	}
	// Should get 200 (CONNECT established)
	response := string(collected)
	if !strings.Contains(response, "200") {
		t.Logf("response: %q", response)
	}
}

// ---------- proxyHTTP: GraphQL read error ----------

func TestProxyHTTPGraphQLReadBodyError(t *testing.T) {
	p := newTestProxy()

	clientConn, proxyClient := net.Pipe()
	upstreamConn, proxyUpstream := net.Pipe()
	defer upstreamConn.Close()

	go p.proxyHTTP(proxyClient, proxyUpstream, "scanner", agent.ModeAdvisory)

	// Send a POST to /graphql with Content-Length but then close before sending body
	go func() {
		fmt.Fprintf(clientConn, "POST /graphql HTTP/1.1\r\nHost: api.github.com\r\nContent-Length: 1000\r\n\r\n")
		// Send partial body then close
		clientConn.Write([]byte(`{"query":"`))
		time.Sleep(50 * time.Millisecond)
		clientConn.Close()
	}()

	go func() {
		buf := make([]byte, 4096)
		for {
			if _, err := upstreamConn.Read(buf); err != nil {
				return
			}
		}
	}()

	// proxyHTTP should handle the read error gracefully
	time.Sleep(500 * time.Millisecond)
}

// ---------- proxyHTTP: git path with upstream write error ----------

func TestProxyHTTPGitPathUpstreamFail(t *testing.T) {
	p := newTestProxy()

	clientConn, proxyClient := net.Pipe()
	upstreamConn, proxyUpstream := net.Pipe()

	go p.proxyHTTP(proxyClient, proxyUpstream, "scanner", agent.ModeIssuesAndPRs)

	// Send a git receive-pack request (write op, allowed at ISSUES_AND_PRS)
	go func() {
		fmt.Fprintf(clientConn, "POST /org/repo.git/git-receive-pack HTTP/1.1\r\nHost: github.com\r\nContent-Length: 5\r\n\r\nhello")
		time.Sleep(100 * time.Millisecond)
		clientConn.Close()
	}()

	go func() {
		// Read request header from upstream
		req, err := http.ReadRequest(bufio.NewReader(upstreamConn))
		if err != nil {
			return
		}
		_ = req
		// Send back some data
		upstreamConn.Write([]byte("git response data"))
		upstreamConn.Close()
	}()

	time.Sleep(300 * time.Millisecond)
}

// ---------- handleAnthropicReroute: forge cert error path ----------

func TestHandleAnthropicRerouteTLSHandshakeFail(t *testing.T) {
	p := newTestProxy()
	route := &InferenceRoute{Backend: "vllm", Endpoint: "http://localhost:1", Model: "test"}

	clientConn, proxyConn := net.Pipe()

	req, _ := http.NewRequest("CONNECT", "http://api.anthropic.com:443", nil)
	req.Host = "api.anthropic.com:443"

	go func() {
		p.handleAnthropicReroute(proxyConn, req, "api.anthropic.com", "scanner", route)
		proxyConn.Close()
	}()

	// Read 200 response
	reader := bufio.NewReader(clientConn)
	line, _ := reader.ReadString('\n')
	if !strings.Contains(line, "200") {
		t.Logf("response: %q", line)
	}
	// Read remaining headers
	for {
		l, _ := reader.ReadString('\n')
		if l == "\r\n" || l == "\n" {
			break
		}
	}

	// Send garbage instead of proper TLS ClientHello — handshake will fail
	clientConn.Write([]byte("not a TLS handshake"))
	time.Sleep(200 * time.Millisecond)
	clientConn.Close()
}

// ---------- handleInferenceRequest: read body failure ----------

func TestHandleInferenceRequestReadBodyFail(t *testing.T) {
	p := newTestProxy()
	route := &InferenceRoute{Backend: "vllm", Endpoint: "http://localhost:1", Model: "test"}

	clientConn, proxyConn := net.Pipe()
	defer clientConn.Close()

	// Create request with a body that errors on read
	req, _ := http.NewRequest("POST", "https://api.anthropic.com/v1/messages", nil)
	req.Body = io.NopCloser(&errorReader{err: fmt.Errorf("read error")})

	go func() {
		p.handleInferenceRequest(proxyConn, req, "scanner", route)
		proxyConn.Close()
	}()

	buf := make([]byte, 4096)
	var collected []byte
	for {
		clientConn.SetReadDeadline(time.Now().Add(2 * time.Second))
		n, err := clientConn.Read(buf)
		if n > 0 {
			collected = append(collected, buf[:n]...)
		}
		if err != nil {
			break
		}
	}
	response := string(collected)
	if !strings.Contains(response, "502") {
		t.Logf("response: %q", response)
	}
}

type errorReader struct {
	err error
}

func (e *errorReader) Read(p []byte) (int, error) {
	return 0, e.err
}

// ---------- handleInferenceRequest: upstream error response ----------

func TestHandleInferenceRequestUpstreamError(t *testing.T) {
	p := newTestProxy()

	mock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte("service down"))
	}))
	defer mock.Close()

	route := &InferenceRoute{Backend: "vllm", Endpoint: mock.URL, Model: "test"}
	body := `{"model":"claude","max_tokens":100,"messages":[{"role":"user","content":"hi"}]}`

	clientConn, proxyConn := net.Pipe()
	defer clientConn.Close()

	req, _ := http.NewRequest("POST", "https://api.anthropic.com/v1/messages", nil)
	req.Body = io.NopCloser(strings.NewReader(body))

	go func() {
		p.handleInferenceRequest(proxyConn, req, "scanner", route)
		proxyConn.Close()
	}()

	// Read response — should be OK (200) with translated content
	// Actually handleInferenceRequest doesn't check status codes — it reads body and translates
	buf := make([]byte, 4096)
	var collected []byte
	for {
		clientConn.SetReadDeadline(time.Now().Add(2 * time.Second))
		n, err := clientConn.Read(buf)
		if n > 0 {
			collected = append(collected, buf[:n]...)
		}
		if err != nil {
			break
		}
	}
	// Response should exist
	if len(collected) == 0 {
		t.Error("expected some response")
	}
}

// ---------- StartInferenceTranslator: response translation error ----------

func TestStartInferenceTranslatorTranslationError(t *testing.T) {
	// Mock backend that returns invalid JSON in non-streaming mode
	mock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("not valid json"))
	}))
	defer mock.Close()

	// Check if port is free
	ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", InferenceTranslatePort))
	if err != nil {
		t.Skipf("port %d not available: %v", InferenceTranslatePort, err)
	}
	ln.Close()

	caCert, caX509, _ := generateCA()
	p := &GitHubProxy{
		caCert:     caCert,
		caX509:     caX509,
		logger:     slog.Default(),
		violations: make(map[string]int),
		certCache:  make(map[string]cachedCert),
		inference:  newInferenceRouter(),
	}
	p.inference.Set("agent-bad-resp", &InferenceRoute{Backend: "vllm", Endpoint: mock.URL, Model: "test"})

	go p.StartInferenceTranslator()
	time.Sleep(200 * time.Millisecond)

	baseURL := fmt.Sprintf("http://127.0.0.1:%d", InferenceTranslatePort)

	// Send request that will get an invalid JSON response from backend
	body := `{"model":"claude","max_tokens":100,"messages":[{"role":"user","content":"hi"}]}`
	req, _ := http.NewRequest("POST", baseURL+"/v1/messages", strings.NewReader(body))
	req.Header.Set("x-api-key", "sk-hive-agent-bad-resp")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	// Should get 502 because response translation fails
	if resp.StatusCode != http.StatusBadGateway {
		t.Errorf("bad response translation should return 502, got %d", resp.StatusCode)
	}
	resp.Body.Close()
}

// ---------- handleConnectDirect: Anthropic reroute TLS handshake fail ----------

func TestHandleConnectDirectAnthropicRerouteForgeCertUsed(t *testing.T) {
	p := newTestProxy()
	p.inference = newInferenceRouter()

	mock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]interface{}{
			"id":      "chatcmpl-test",
			"choices": []map[string]interface{}{{"message": map[string]string{"content": "rerouted"}, "finish_reason": "stop"}},
			"usage":   map[string]int{"prompt_tokens": 1, "completion_tokens": 1},
		})
	}))
	defer mock.Close()

	route := &InferenceRoute{Backend: "vllm", Endpoint: mock.URL, Model: "test"}
	p.SetInferenceRoute("agent-reroute", route)

	clientConn, proxyClient := net.Pipe()
	defer clientConn.Close()

	req, _ := http.NewRequest("CONNECT", "http://api.anthropic.com:443", nil)
	req.Host = "api.anthropic.com:443"
	req.Header.Set("Proxy-Authorization", "hive agent-reroute")

	go func() {
		p.handleConnectDirect(proxyClient, req)
		proxyClient.Close()
	}()

	// Read 200
	reader := bufio.NewReader(clientConn)
	line, _ := reader.ReadString('\n')
	if !strings.Contains(line, "200") {
		t.Logf("expected 200, got %q", line)
	}
}

// ---------- handleTransparentTLS: GitHub host with mode file ----------

func TestHandleTransparentTLSGitHubWithMode(t *testing.T) {
	p := newTestProxy()
	p.uidMap = nil // No UID map

	clientHello := buildClientHello("github.com")
	peeked := clientHello[:1]

	clientConn, proxyConn := net.Pipe()

	go func() {
		p.handleTransparentTLS(proxyConn, peeked)
		proxyConn.Close()
	}()

	clientConn.Write(clientHello[1:])
	// It will try to forge cert and TLS handshake, then dial github.com:443
	// All will complete or fail gracefully
	time.Sleep(500 * time.Millisecond)
	clientConn.Close()
}

// ---------- proxyHTTP: response write succeeds then loops ----------

func TestProxyHTTPSuccessfulRoundtrip(t *testing.T) {
	p := newTestProxy()

	clientConn, proxyClient := net.Pipe()
	upstreamConn, proxyUpstream := net.Pipe()

	go p.proxyHTTP(proxyClient, proxyUpstream, "scanner", agent.ModeIssuesAndPRs)

	// Request 1: allowed POST (create issue)
	body := `{"title":"test"}`
	go func() {
		fmt.Fprintf(clientConn, "POST /repos/org/repo/issues HTTP/1.1\r\nHost: api.github.com\r\nContent-Length: %d\r\n\r\n%s", len(body), body)
		// After getting response, close
		time.Sleep(300 * time.Millisecond)
		clientConn.Close()
	}()

	go func() {
		req, err := http.ReadRequest(bufio.NewReader(upstreamConn))
		if err != nil {
			return
		}
		io.ReadAll(req.Body)
		resp := &http.Response{
			StatusCode:    201,
			Proto:         "HTTP/1.1",
			ProtoMajor:    1,
			ProtoMinor:    1,
			Header:        http.Header{"Content-Type": []string{"application/json"}},
			Body:          io.NopCloser(strings.NewReader(`{"id":42}`)),
			ContentLength: 9,
		}
		resp.Write(upstreamConn)
		upstreamConn.Close()
	}()

	reader := bufio.NewReader(clientConn)
	resp, err := http.ReadResponse(reader, nil)
	if err != nil {
		t.Logf("read response: %v (may be expected if conn closed)", err)
		return
	}
	if resp.StatusCode != 201 {
		t.Errorf("status = %d, want 201", resp.StatusCode)
	}
}
