package server

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

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

func TestAdminExtractConfigUnsupportedTypes(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	for _, mediaType := range []string{"text", "bogus"} {
		for _, method := range []string{http.MethodGet, http.MethodPut} {
			resp := doExtractConfigRequest(t, ts, method, rt.tenantID, mediaType, `{"enabled":false}`, legacyExtractHeaders())
			if resp.StatusCode != http.StatusBadRequest {
				t.Fatalf("%s %s: status = %d, want 400", method, mediaType, resp.StatusCode)
			}
			out := decodeExtractConfigBody(t, resp)
			want := fmt.Sprintf("extract config for %q is not supported", mediaType)
			if out["error"] != want {
				t.Fatalf("%s %s: error = %v, want %q", method, mediaType, out["error"], want)
			}
		}
	}
}

func TestAdminExtractConfigSupportsAudioAndVideo(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	ts := httptest.NewServer(rt.server)
	defer ts.Close()
	for _, mediaType := range []string{"audio", "video"} {
		resp := doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, mediaType, `{"enabled":false}`, legacyExtractHeaders())
		if resp.StatusCode != http.StatusOK {
			t.Fatalf("PUT %s: status=%d, want 200", mediaType, resp.StatusCode)
		}
		_ = decodeExtractConfigBody(t, resp)
	}
}

func TestAdminExtractConfigUnknownTenant(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	getResp := doExtractConfigRequest(t, ts, http.MethodGet, "tenant-unknown", "image", "", legacyExtractHeaders())
	if getResp.StatusCode != http.StatusNotFound {
		t.Fatalf("GET status = %d, want 404", getResp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, getResp)
	putResp := doExtractConfigRequest(t, ts, http.MethodPut, "tenant-unknown", "image", `{"enabled":false}`, legacyExtractHeaders())
	if putResp.StatusCode != http.StatusNotFound {
		t.Fatalf("PUT status = %d, want 404", putResp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, putResp)
}

func TestAdminExtractConfigCrossOrgRejectedNative(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", nil)
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{
		OrganizationID: "org-2", Role: tenant.TiDBCloudRoleOrgOwner,
	}}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	getResp := doExtractConfigRequest(t, ts, http.MethodGet, rt.tenantID, "image", "", legacyExtractHeaders())
	if getResp.StatusCode != http.StatusForbidden {
		t.Fatalf("GET status = %d, want 403", getResp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, getResp)
	putResp := doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, "image", `{"enabled":false}`, legacyExtractHeaders())
	if putResp.StatusCode != http.StatusForbidden {
		t.Fatalf("PUT status = %d, want 403", putResp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, putResp)
}

func TestAdminExtractConfigSharedProvider(t *testing.T) {
	rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNativeShared, "", nil)
	bindSharedExtractTenant(t, rt, "org-shared")
	// Identities resolve in call order; the RBAC cache keys by credentials, so
	// the cross-org phase uses a distinct credential pair.
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{
		{OrganizationID: "org-shared", Role: tenant.TiDBCloudRoleOrgOwner},
		{OrganizationID: "org-2", Role: tenant.TiDBCloudRoleOrgOwner},
		{OrganizationID: "org-2", Role: tenant.TiDBCloudRoleOrgOwner},
	}
	ts := httptest.NewServer(rt.server)
	defer ts.Close()

	// Happy path: identity org matches the shared pool organization.
	getResp := doExtractConfigRequest(t, ts, http.MethodGet, rt.tenantID, "image", "", legacyExtractHeaders())
	if getResp.StatusCode != http.StatusOK {
		t.Fatalf("shared GET status = %d, want 200", getResp.StatusCode)
	}
	out := decodeExtractConfigBody(t, getResp)
	if out["enabled"] != false || out["source"] != "none" {
		t.Fatalf("shared GET response = %v", out)
	}
	putResp := doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, "image",
		`{"enabled":true,"api_base":"https://api.example.com","api_key":"sk-shared-key","model":"model-1"}`, legacyExtractHeaders())
	if putResp.StatusCode != http.StatusOK {
		t.Fatalf("shared PUT status = %d, want 200", putResp.StatusCode)
	}
	out = decodeExtractConfigBody(t, putResp)
	if out["enabled"] != true || out["source"] != "custom" || out["api_key"] != maskAPIKey("sk-shared-key") {
		t.Fatalf("shared PUT response = %v", out)
	}

	// Cross-org: identity org does not own the shared pool.
	rt.prov.iamIdentities = []*tenant.TiDBCloudAPIKeyIdentity{{
		OrganizationID: "org-2", Role: tenant.TiDBCloudRoleOrgOwner,
	}}
	crossOrgHeaders := map[string]string{quotaPublicKeyHeader: "public-2", quotaPrivateKeyHeader: "private-2"}
	getResp = doExtractConfigRequest(t, ts, http.MethodGet, rt.tenantID, "image", "", crossOrgHeaders)
	if getResp.StatusCode != http.StatusForbidden {
		t.Fatalf("cross-org shared GET status = %d, want 403", getResp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, getResp)
	putResp = doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, "image", `{"enabled":false}`, crossOrgHeaders)
	if putResp.StatusCode != http.StatusForbidden {
		t.Fatalf("cross-org shared PUT status = %d, want 403", putResp.StatusCode)
	}
	_ = decodeExtractConfigBody(t, putResp)
}

func TestAdminExtractConfigDispatcherMode(t *testing.T) {
	newDispatcher := func(account *innerAuthAccountFake) *innerauth.Dispatcher {
		return innerauth.NewDispatcher(account, innerAuthDefaultsFake{
			credentials: tenant.CredentialProvisionRequest{PublicKey: "server-public", PrivateKey: "server-private"},
			ok:          true,
		}, time.Hour)
	}
	callerHeaders := map[string]string{quotaPublicKeyHeader: "caller-public", quotaPrivateKeyHeader: "caller-private"}

	t.Run("GET and PUT authorize with extract-config resource URIs", func(t *testing.T) {
		account := &innerAuthAccountFake{apiKeyResult: innerauth.APIKeyAuthnResult{PublicKey: "caller-public", OrganizationID: 60205}}
		rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "60205", newDispatcher(account))
		ts := httptest.NewServer(rt.server)
		defer ts.Close()

		getResp := doExtractConfigRequest(t, ts, http.MethodGet, rt.tenantID, "image", "", callerHeaders)
		if getResp.StatusCode != http.StatusOK {
			t.Fatalf("GET status = %d, want 200", getResp.StatusCode)
		}
		out := decodeExtractConfigBody(t, getResp)
		if out["source"] != "none" {
			t.Fatalf("GET response = %v", out)
		}
		if account.requestURI != "/drive9/orgs/60205/extract-config/view" ||
			account.resourceURI != "/drive9/orgs/{org_id}/extract-config/view" {
			t.Fatalf("GET authz URIs = (%q, %q)", account.requestURI, account.resourceURI)
		}

		putResp := doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, "image",
			`{"enabled":true,"api_base":"https://api.example.com","api_key":"sk-dispatcher","model":"model-1"}`, callerHeaders)
		if putResp.StatusCode != http.StatusOK {
			t.Fatalf("PUT status = %d, want 200", putResp.StatusCode)
		}
		out = decodeExtractConfigBody(t, putResp)
		if out["enabled"] != true || out["api_key"] != maskAPIKey("sk-dispatcher") {
			t.Fatalf("PUT response = %v", out)
		}
		if account.requestURI != "/drive9/orgs/60205/extract-config/modify" ||
			account.resourceURI != "/drive9/orgs/{org_id}/extract-config/modify" {
			t.Fatalf("PUT authz URIs = (%q, %q)", account.requestURI, account.resourceURI)
		}
	})

	t.Run("cross-organization subject is forbidden", func(t *testing.T) {
		account := &innerAuthAccountFake{apiKeyResult: innerauth.APIKeyAuthnResult{PublicKey: "caller-public", OrganizationID: 60205}}
		// Tenant is bound to org-1 while the authenticated subject is org 60205.
		rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "org-1", newDispatcher(account))
		ts := httptest.NewServer(rt.server)
		defer ts.Close()

		getResp := doExtractConfigRequest(t, ts, http.MethodGet, rt.tenantID, "image", "", callerHeaders)
		if getResp.StatusCode != http.StatusForbidden {
			t.Fatalf("GET status = %d, want 403", getResp.StatusCode)
		}
		_ = decodeExtractConfigBody(t, getResp)
		putResp := doExtractConfigRequest(t, ts, http.MethodPut, rt.tenantID, "image", `{"enabled":false}`, callerHeaders)
		if putResp.StatusCode != http.StatusForbidden {
			t.Fatalf("PUT status = %d, want 403", putResp.StatusCode)
		}
		_ = decodeExtractConfigBody(t, putResp)
	})

	t.Run("middleware authorization denial short-circuits", func(t *testing.T) {
		account := &innerAuthAccountFake{
			apiKeyResult: innerauth.APIKeyAuthnResult{PublicKey: "caller-public", OrganizationID: 60205},
			authzErr:     innerauth.ErrForbidden,
		}
		rt := newExtractConfigRuntime(t, tenant.ProviderTiDBCloudNative, "60205", newDispatcher(account))
		ts := httptest.NewServer(rt.server)
		defer ts.Close()

		resp := doExtractConfigRequest(t, ts, http.MethodGet, rt.tenantID, "image", "", callerHeaders)
		if resp.StatusCode != http.StatusForbidden {
			t.Fatalf("GET status = %d, want 403", resp.StatusCode)
		}
		_ = decodeExtractConfigBody(t, resp)
	})
}