package webagent

import (
	"context"
	"encoding/json"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api/dto"
	"github.com/tidbcloud/AgentTeam/internal/api/testkit/httputil"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
)

func TestModelProfileCreateHandlerCreatesMultipleProfiles(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 5, 24, 9, 0, 0, 0, time.UTC)
	fake := newFakeStore()
	fake.createdProfiles = []agent.ModelProfile{
		{
			ID:               "ump_1",
			ProviderID:       "umpvd_1",
			UserID:           "usr_1",
			DisplayName:      "OpenAI compatible",
			Model:            "gpt-5.4",
			BaseURL:          "https://api.example.com/v1",
			WireAPI:          "openai",
			APIKeyConfigured: true,
			CreatedAt:        now,
			UpdatedAt:        now,
		},
		{
			ID:               "ump_2",
			ProviderID:       "umpvd_1",
			UserID:           "usr_1",
			DisplayName:      "OpenAI compatible",
			Model:            "gpt-5.4-mini",
			BaseURL:          "https://api.example.com/v1",
			WireAPI:          "openai",
			APIKeyConfigured: true,
			CreatedAt:        now,
			UpdatedAt:        now,
		},
	}
	router := newWebAgentRouter(Dependencies{AgentStore: fake, ReviewStore: fake})

	resp := httputil.DoJSON(t, router, http.MethodPost, "/api/v1/me/model-profiles", map[string]any{
		"displayName": " OpenAI compatible ",
		"models":      []string{" gpt-5.4 ", "gpt-5.4", "gpt-5.4-mini"},
		"baseUrl":     " https://api.example.com/v1 ",
		"apiKey":      " sk-test ",
		"wireApi":     " openai ",
	})
	require.Equal(t, http.StatusOK, resp.Code)

	require.Equal(t, []string{"gpt-5.4", "gpt-5.4-mini"}, fake.lastCreateProfile.Models)
	require.Equal(t, " OpenAI compatible ", fake.lastCreateProfile.DisplayName)
	require.Equal(t, " https://api.example.com/v1 ", fake.lastCreateProfile.BaseURL)
	require.Equal(t, " sk-test ", fake.lastCreateProfile.APIKey)
	require.Equal(t, " openai ", fake.lastCreateProfile.WireAPI)

	var out []dto.ModelProfile
	require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &out))
	require.Len(t, out, 2)
	require.Equal(t, "ump_1", out[0].ID)
	require.Equal(t, "umpvd_1", out[0].ProviderID)
	require.Equal(t, "gpt-5.4", out[0].Model)
	require.True(t, out[0].APIKeyConfigured)
	require.Equal(t, "ump_2", out[1].ID)
	require.Equal(t, "umpvd_1", out[1].ProviderID)
	require.Equal(t, "gpt-5.4-mini", out[1].Model)
}

func TestModelProfileUpdateHandlerUpdatesProviderProfiles(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 5, 24, 9, 0, 0, 0, time.UTC)
	fake := newFakeStore()
	fake.updatedProfiles = []agent.ModelProfile{
		{
			ID:               "ump_1",
			ProviderID:       "umpvd_1",
			UserID:           "usr_1",
			DisplayName:      "Updated Provider",
			Model:            "gpt-5.5",
			BaseURL:          "https://api.example.com/v2",
			WireAPI:          "responses",
			APIKeyConfigured: true,
			CreatedAt:        now,
			UpdatedAt:        now,
		},
	}
	router := newWebAgentRouter(Dependencies{AgentStore: fake, ReviewStore: fake})

	resp := httputil.DoJSON(t, router, http.MethodPatch, "/api/v1/me/model-profiles/ump_1", map[string]any{
		"displayName": " Updated Provider ",
		"models":      []string{" gpt-5.5 "},
		"baseUrl":     " https://api.example.com/v2 ",
		"apiKey":      " sk-updated ",
		"wireApi":     " responses ",
	})
	require.Equal(t, http.StatusOK, resp.Code)

	require.Equal(t, "ump_1", fake.lastUpdateProfileID)
	require.Equal(t, []string{"gpt-5.5"}, fake.lastUpdateProfile.Models)
	require.Equal(t, " Updated Provider ", fake.lastUpdateProfile.DisplayName)
	require.Equal(t, " https://api.example.com/v2 ", fake.lastUpdateProfile.BaseURL)
	require.Equal(t, " sk-updated ", fake.lastUpdateProfile.APIKey)
	require.Equal(t, " responses ", fake.lastUpdateProfile.WireAPI)

	var out []dto.ModelProfile
	require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &out))
	require.Len(t, out, 1)
	require.Equal(t, "ump_1", out[0].ID)
	require.Equal(t, "Updated Provider", out[0].DisplayName)
	require.Equal(t, "gpt-5.5", out[0].Model)
}

func TestModelProfileGetHandlerReturnsAPIKey(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 5, 24, 9, 0, 0, 0, time.UTC)
	fake := newFakeStore()
	fake.modelProfile = agent.ModelProfile{
		ID:               "ump_1",
		ProviderID:       "umpvd_1",
		UserID:           "usr_1",
		DisplayName:      "OpenAI compatible",
		Model:            "gpt-5.4",
		BaseURL:          "https://api.example.com/v1",
		APIKey:           "sk-test",
		WireAPI:          "openai",
		APIKeyConfigured: true,
		CreatedAt:        now,
		UpdatedAt:        now,
	}
	router := newWebAgentRouter(Dependencies{AgentStore: fake, ReviewStore: fake})

	resp := httputil.DoJSON(t, router, http.MethodGet, "/api/v1/me/model-profiles/ump_1", nil)
	require.Equal(t, http.StatusOK, resp.Code)

	var out dto.ModelProfile
	require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &out))
	require.Equal(t, "ump_1", out.ID)
	require.Equal(t, "sk-test", out.APIKey)
	require.True(t, out.APIKeyConfigured)
}

func TestModelProfileConnectionTestHandlerSendsWireAPIRequests(t *testing.T) {
	withPublicModelProfileLookup(t)
	withModelProfileDialToLoopback(t)

	testCases := []struct {
		name       string
		wireAPI    string
		baseSuffix string
		wantPath   string
	}{
		{
			name:       "openai chat completions",
			wireAPI:    "openai",
			baseSuffix: "/v1",
			wantPath:   "/v1/chat/completions",
		},
		{
			name:       "responses",
			wireAPI:    "responses",
			baseSuffix: "/v1",
			wantPath:   "/v1/responses",
		},
		{
			name:     "anthropic messages",
			wireAPI:  "anthropic",
			wantPath: "/v1/messages",
		},
	}

	for _, testCase := range testCases {
		t.Run(testCase.name, func(t *testing.T) {
			var receivedPath string
			var receivedHeaders http.Header
			var receivedBody map[string]any
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				receivedPath = r.URL.Path
				receivedHeaders = r.Header.Clone()
				require.NoError(t, json.NewDecoder(r.Body).Decode(&receivedBody))
				w.Header().Set("Content-Type", "application/json")
				_, _ = w.Write([]byte(`{"id":"probe_ok","status":"completed"}`))
			}))
			defer upstream.Close()

			fake := newFakeStore()
			router := newWebAgentRouter(Dependencies{AgentStore: fake, ReviewStore: fake})
			resp := httputil.DoJSON(t, router, http.MethodPost, "/api/v1/me/model-profiles:testConnection", map[string]any{
				"baseUrl": modelProfileTestBaseURL(t, upstream.URL, testCase.baseSuffix),
				"apiKey":  "sk-test",
				"wireApi": testCase.wireAPI,
				"model":   "gpt-test",
			})
			require.Equal(t, http.StatusOK, resp.Code)
			require.JSONEq(t, `{}`, resp.Body.String())

			require.Equal(t, testCase.wantPath, receivedPath)
			require.Equal(t, "application/json", receivedHeaders.Get("Content-Type"))
			require.Equal(t, "gpt-test", receivedBody["model"])
			switch testCase.wireAPI {
			case "anthropic":
				require.Equal(t, "sk-test", receivedHeaders.Get("x-api-key"))
				require.Equal(t, "2023-06-01", receivedHeaders.Get("anthropic-version"))
				require.Equal(t, float64(1), receivedBody["max_tokens"])
				requireProbeMessageBody(t, receivedBody)
			case "responses":
				require.Equal(t, "Bearer sk-test", receivedHeaders.Get("Authorization"))
				require.Equal(t, "Reply with OK.", receivedBody["input"])
				require.Equal(t, float64(1), receivedBody["max_output_tokens"])
				require.Equal(t, false, receivedBody["stream"])
			case "openai":
				require.Equal(t, "Bearer sk-test", receivedHeaders.Get("Authorization"))
				require.Equal(t, float64(1), receivedBody["max_tokens"])
				require.Equal(t, false, receivedBody["stream"])
				requireProbeMessageBody(t, receivedBody)
			default:
				t.Fatalf("unhandled wire API %q", testCase.wireAPI)
			}
		})
	}
}

func TestModelProfileConnectionTestHandlerRejectsUpstreamAuthFailure(t *testing.T) {
	withPublicModelProfileLookup(t)
	withModelProfileDialToLoopback(t)

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "secret upstream detail", http.StatusUnauthorized)
	}))
	defer upstream.Close()

	fake := newFakeStore()
	router := newWebAgentRouter(Dependencies{AgentStore: fake, ReviewStore: fake})
	resp := httputil.DoJSON(t, router, http.MethodPost, "/api/v1/me/model-profiles:testConnection", map[string]any{
		"baseUrl": modelProfileTestBaseURL(t, upstream.URL, "/v1"),
		"apiKey":  "sk-test",
		"wireApi": "openai",
		"model":   "gpt-test",
	})
	require.Equal(t, http.StatusBadGateway, resp.Code)
	require.Contains(t, resp.Body.String(), "API key was rejected")
	require.NotContains(t, resp.Body.String(), "secret upstream detail")
}

func TestModelProfileConnectionTestHandlerDoesNotFollowRedirects(t *testing.T) {
	withPublicModelProfileLookup(t)
	withModelProfileDialToLoopback(t)

	redirectTargetHit := false
	redirectTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		redirectTargetHit = true
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"unexpected"}`))
	}))
	defer redirectTarget.Close()

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Location", redirectTarget.URL+"/v1/chat/completions")
		w.WriteHeader(http.StatusTemporaryRedirect)
		_, _ = w.Write([]byte("redirect detail"))
	}))
	defer upstream.Close()

	fake := newFakeStore()
	router := newWebAgentRouter(Dependencies{AgentStore: fake, ReviewStore: fake})
	resp := httputil.DoJSON(t, router, http.MethodPost, "/api/v1/me/model-profiles:testConnection", map[string]any{
		"baseUrl": modelProfileTestBaseURL(t, upstream.URL, "/v1"),
		"apiKey":  "sk-test",
		"wireApi": "openai",
		"model":   "gpt-test",
	})
	require.Equal(t, http.StatusBadGateway, resp.Code)
	require.Contains(t, resp.Body.String(), "test endpoint returned HTTP 307")
	require.False(t, redirectTargetHit)
}

func TestModelProfileConnectionTestHandlerRejectsLoopbackHostAtDial(t *testing.T) {
	lookupCount := 0
	originalLookup := modelProfileLookupIP
	modelProfileLookupIP = func(context.Context, string, string) ([]net.IP, error) {
		lookupCount++
		return []net.IP{net.ParseIP("127.0.0.1")}, nil
	}
	t.Cleanup(func() {
		modelProfileLookupIP = originalLookup
	})

	dialCalled := false
	originalDial := modelProfileDialContext
	modelProfileDialContext = func(context.Context, string, string) (net.Conn, error) {
		dialCalled = true
		return nil, context.Canceled
	}
	t.Cleanup(func() {
		modelProfileDialContext = originalDial
	})

	fake := newFakeStore()
	router := newWebAgentRouter(Dependencies{AgentStore: fake, ReviewStore: fake})
	resp := httputil.DoJSON(t, router, http.MethodPost, "/api/v1/me/model-profiles:testConnection", map[string]any{
		"baseUrl": "http://rebind.example.test/v1",
		"apiKey":  "sk-test",
		"wireApi": "openai",
		"model":   "gpt-test",
	})
	require.Equal(t, http.StatusBadRequest, resp.Code)
	require.Contains(t, resp.Body.String(), "Base URL host is not allowed")
	require.Equal(t, 1, lookupCount)
	require.False(t, dialCalled)
}

func TestModelProfileHostIPsAllowPrivateGateways(t *testing.T) {
	for _, address := range []string{
		"10.20.30.40",
		"172.20.30.40",
		"192.168.40.106",
		"fd12:3456:789a::1",
	} {
		t.Run(address, func(t *testing.T) {
			ips, err := modelProfileHostIPs(context.Background(), address)
			require.NoError(t, err)
			require.Equal(t, []net.IP{normalizeModelProfileIP(net.ParseIP(address))}, ips)
		})
	}
}

func TestModelProfileHostIPsStillRejectSensitiveLocalAddresses(t *testing.T) {
	for _, address := range []string{
		"127.0.0.1",
		"169.254.169.254",
		"100.64.0.1",
		"::1",
		"fe80::1",
	} {
		t.Run(address, func(t *testing.T) {
			_, err := modelProfileHostIPs(context.Background(), address)
			require.Error(t, err)
			require.Contains(t, err.Error(), "Base URL host is not allowed")
		})
	}
}

func requireProbeMessageBody(t *testing.T, body map[string]any) {
	t.Helper()
	messages, ok := body["messages"].([]any)
	require.True(t, ok)
	require.Len(t, messages, 1)
	message, ok := messages[0].(map[string]any)
	require.True(t, ok)
	require.Equal(t, "user", message["role"])
	require.Equal(t, "Reply with OK.", message["content"])
}

func withPublicModelProfileLookup(t *testing.T) {
	t.Helper()
	original := modelProfileLookupIP
	modelProfileLookupIP = func(context.Context, string, string) ([]net.IP, error) {
		return []net.IP{net.ParseIP("93.184.216.34")}, nil
	}
	t.Cleanup(func() {
		modelProfileLookupIP = original
	})
}

func withModelProfileDialToLoopback(t *testing.T) {
	t.Helper()
	original := modelProfileDialContext
	dialer := &net.Dialer{}
	modelProfileDialContext = func(ctx context.Context, network string, address string) (net.Conn, error) {
		_, port, err := net.SplitHostPort(address)
		if err != nil {
			return nil, err
		}
		return dialer.DialContext(ctx, network, net.JoinHostPort("127.0.0.1", port))
	}
	t.Cleanup(func() {
		modelProfileDialContext = original
	})
}

func modelProfileTestBaseURL(t *testing.T, serverURL string, suffix string) string {
	t.Helper()
	parsed, err := url.Parse(serverURL)
	require.NoError(t, err)
	_, port, err := net.SplitHostPort(parsed.Host)
	require.NoError(t, err)
	return parsed.Scheme + "://model-profile-test.example:" + port + suffix
}