package server

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/tenant"
)

func TestQuotaGetExposesLLMLimitsUsageAndDedicatedSpendingLimit(t *testing.T) {
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
	ctx := context.Background()
	if err := rt.meta.SetQuotaConfig(ctx, &meta.QuotaConfig{
		TenantID: rt.tenantID, MaxStorageBytes: 123 * quotaStorageSizeBytes,
		MaxFileSizeBytes: 12 * quotaStorageSizeBytes, MaxFileCount: 34,
		MaxMediaLLMFiles: 56, MaxVideoLLMFiles: 78, MaxMonthlyCostMC: 900,
	}); err != nil {
		t.Fatal(err)
	}
	if err := rt.meta.EnsureQuotaUsageRow(ctx, rt.tenantID); err != nil {
		t.Fatal(err)
	}
	if _, err := rt.meta.DB().ExecContext(ctx, `UPDATE tenant_quota_usage
		SET storage_bytes = 321, reserved_bytes = 11, file_count = 9,
			media_file_count = 7, video_file_count = 5
		WHERE tenant_id = ?`, rt.tenantID); err != nil {
		t.Fatal(err)
	}

	ts := httptest.NewServer(rt.server)
	defer ts.Close()
	resp := getQuota(t, ts.URL, rt.tenantID, "public-1", "private-1", "")
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		t.Fatalf("status = %d, want 200: %s", resp.StatusCode, body)
	}
	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	var out quotaResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		t.Fatal(err)
	}
	if out.Config.MaxMediaLLMFiles != 56 || out.Config.MaxVideoLLMFiles != 78 {
		t.Fatalf("LLM limits = %+v", out.Config)
	}
	if out.Usage.MediaFileCount != 7 || out.Usage.VideoFileCount != 5 {
		t.Fatalf("LLM usage = %+v", out.Usage)
	}
	if !strings.Contains(string(raw), `"tidbcloud_spending_limit":null`) {
		t.Fatalf("response omitted dedicated spending limit: %s", raw)
	}
	if strings.Contains(string(raw), "max_monthly_cost_mc") {
		t.Fatalf("response exposed internal monthly cost limit: %s", raw)
	}
}

func TestQuotaSetLLMLimitsRequiresEnabledTenantExtractConfig(t *testing.T) {
	tests := []struct {
		name      string
		field     string
		mediaType string
		want      func(*meta.QuotaConfig) int64
	}{
		{
			name: "media", field: "max_media_llm_files", mediaType: meta.ExtractMediaTypeImage,
			want: func(cfg *meta.QuotaConfig) int64 { return cfg.MaxMediaLLMFiles },
		},
		{
			name: "video", field: "max_video_llm_files", mediaType: meta.ExtractMediaTypeVideo,
			want: func(cfg *meta.QuotaConfig) int64 { return cfg.MaxVideoLLMFiles },
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
			ts := httptest.NewServer(rt.server)
			defer ts.Close()
			postLimit := func(value int64) *http.Response {
				t.Helper()
				return postJSON(t, ts.URL+"/v1/quota", map[string]any{
					"tenant_id": rt.tenantID, "public_key": "public-1", "private_key": "private-1",
					tt.field: value,
				}, "")
			}

			resp := postLimit(123)
			body, _ := io.ReadAll(resp.Body)
			_ = resp.Body.Close()
			if resp.StatusCode != http.StatusConflict || !strings.Contains(string(body), "enabled tenant extract config") {
				t.Fatalf("default config response = %d %s, want 409 extract-config error", resp.StatusCode, body)
			}

			if err := rt.meta.UpsertTenantExtractConfig(context.Background(), &meta.TenantExtractConfig{
				TenantID: rt.tenantID, MediaType: tt.mediaType, Enabled: false,
			}); err != nil {
				t.Fatal(err)
			}
			resp = postLimit(123)
			body, _ = io.ReadAll(resp.Body)
			_ = resp.Body.Close()
			if resp.StatusCode != http.StatusConflict || !strings.Contains(string(body), "enabled tenant extract config") {
				t.Fatalf("disabled config response = %d %s, want 409 extract-config error", resp.StatusCode, body)
			}

			if err := rt.meta.UpsertTenantExtractConfig(context.Background(), &meta.TenantExtractConfig{
				TenantID: rt.tenantID, MediaType: tt.mediaType, Enabled: true,
				APIBase: "https://extract.example.com/v1", APIKeyCipher: []byte("cipher"), Model: "model-1",
			}); err != nil {
				t.Fatal(err)
			}
			resp = postLimit(123)
			body, _ = io.ReadAll(resp.Body)
			_ = resp.Body.Close()
			if resp.StatusCode != http.StatusOK {
				t.Fatalf("enabled config response = %d %s, want 200", resp.StatusCode, body)
			}
			cfg, err := rt.meta.GetQuotaConfig(context.Background(), rt.tenantID)
			if err != nil {
				t.Fatal(err)
			}
			if got := tt.want(cfg); got != 123 {
				t.Fatalf("stored %s = %d, want 123", tt.field, got)
			}
		})
	}
}

func TestQuotaSetRejectsLLMLimitBeforeDedicatedSpendingLimitPatch(t *testing.T) {
	rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
	spendingLimit := int64(12_000)
	mediaLimit := int64(123)
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	resp := postJSON(t, ts.URL+"/v1/quota", map[string]any{
		"tenant_id": rt.tenantID, "public_key": "public-1", "private_key": "private-1",
		"tidbcloud_spending_limit": spendingLimit,
		"max_media_llm_files":      mediaLimit,
	}, "")
	defer func() { _ = resp.Body.Close() }()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusConflict || !strings.Contains(string(body), "enabled tenant extract config") {
		t.Fatalf("response = %d %s, want 409 extract-config error", resp.StatusCode, body)
	}
	if got := rt.prov.markCalls.Load() + rt.prov.updateCalls.Load(); got != 0 {
		t.Fatalf("dedicated cloud mutation calls = %d, want 0", got)
	}
	cfg, err := rt.meta.GetQuotaConfig(context.Background(), rt.tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.TiDBCloudSpendingLimit != nil {
		t.Fatalf("persisted spending limit = %#v, want unchanged nil", cfg.TiDBCloudSpendingLimit)
	}
}

func TestAdminTenantQuotaSetLLMLimitsRequiresEnabledTenantExtractConfig(t *testing.T) {
	tests := []struct {
		name      string
		field     string
		mediaType string
		want      func(*meta.QuotaConfig) int64
	}{
		{
			name: "media", field: "max_media_llm_files", mediaType: meta.ExtractMediaTypeImage,
			want: func(cfg *meta.QuotaConfig) int64 { return cfg.MaxMediaLLMFiles },
		},
		{
			name: "video", field: "max_video_llm_files", mediaType: meta.ExtractMediaTypeVideo,
			want: func(cfg *meta.QuotaConfig) int64 { return cfg.MaxVideoLLMFiles },
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			rt := newQuotaRuntime(t, tenant.ProviderTiDBCloudNative)
			rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{
				OrganizationID: "org-1", Role: tenant.TiDBCloudRoleOrgOwner,
			}}
			ts := httptest.NewServer(rt.server)
			defer ts.Close()
			postLimit := func(value int64) *http.Response {
				t.Helper()
				return postJSON(t, ts.URL+"/v1/admin/tenants/"+rt.tenantID+"/quota", map[string]any{
					"public_key": "public-1", "private_key": "private-1", tt.field: value,
				}, "")
			}

			resp := postLimit(234)
			body, _ := io.ReadAll(resp.Body)
			_ = resp.Body.Close()
			if resp.StatusCode != http.StatusConflict || !strings.Contains(string(body), "enabled tenant extract config") {
				t.Fatalf("default config response = %d %s, want 409 extract-config error", resp.StatusCode, body)
			}

			if err := rt.meta.UpsertTenantExtractConfig(context.Background(), &meta.TenantExtractConfig{
				TenantID: rt.tenantID, MediaType: tt.mediaType, Enabled: false,
			}); err != nil {
				t.Fatal(err)
			}
			resp = postLimit(234)
			body, _ = io.ReadAll(resp.Body)
			_ = resp.Body.Close()
			if resp.StatusCode != http.StatusConflict || !strings.Contains(string(body), "enabled tenant extract config") {
				t.Fatalf("disabled config response = %d %s, want 409 extract-config error", resp.StatusCode, body)
			}

			if err := rt.meta.UpsertTenantExtractConfig(context.Background(), &meta.TenantExtractConfig{
				TenantID: rt.tenantID, MediaType: tt.mediaType, Enabled: true,
				APIBase: "https://extract.example.com/v1", APIKeyCipher: []byte("cipher"), Model: "model-1",
			}); err != nil {
				t.Fatal(err)
			}
			resp = postLimit(234)
			body, _ = io.ReadAll(resp.Body)
			_ = resp.Body.Close()
			if resp.StatusCode != http.StatusOK {
				t.Fatalf("enabled config response = %d %s, want 200", resp.StatusCode, body)
			}
			cfg, err := rt.meta.GetQuotaConfig(context.Background(), rt.tenantID)
			if err != nil {
				t.Fatal(err)
			}
			if got := tt.want(cfg); got != 234 {
				t.Fatalf("stored %s = %d, want 234", tt.field, got)
			}
		})
	}
}