package organization

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/core/secrets"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

func (txs TxStore) hasOrganizationProxyKeySigningKey() bool {
	return len(txs.feishuCredentialEncryptionKey) == 32
}

func (txs TxStore) CreateInternalWorkspaceProxyKey(ctx context.Context, workspaceID string, expiresAt time.Time) (string, error) {
	if !txs.hasOrganizationProxyKeySigningKey() {
		return "", badRequest("organization proxy signing key is not configured on server")
	}
	_ = ctx
	_ = expiresAt
	payload := model.SignedProviderProxyKey{
		WorkspaceID:   strings.TrimSpace(workspaceID),
		ExpiresAtUnix: 0,
	}
	if payload.WorkspaceID == "" {
		return "", badRequest("workspace is required")
	}
	rawPayload, err := model.MarshalJSONObject("signed provider proxy key", payload)
	if err != nil {
		return "", err
	}
	encodedPayload := base64.RawURLEncoding.EncodeToString(rawPayload)
	mac := hmac.New(sha256.New, txs.feishuCredentialEncryptionKey)
	_, _ = mac.Write([]byte(encodedPayload))
	signature := hex.EncodeToString(mac.Sum(nil))
	return "lpi_" + encodedPayload + "." + signature, nil
}

func (txs TxStore) AuthenticateOrganizationLLMProxyKey(ctx context.Context, token string) (OrganizationProviderProxyAccess, error) {
	access, err := txs.AuthenticatePersistentOrganizationLLMProxyKey(ctx, token)
	if err == nil {
		return access, nil
	}
	if !errors.Is(err, ErrNotFound) {
		return OrganizationProviderProxyAccess{}, err
	}
	return txs.AuthenticateSignedWorkspaceLLMProxyKey(ctx, token)
}

func (txs TxStore) AuthenticatePersistentOrganizationLLMProxyKey(ctx context.Context, token string) (OrganizationProviderProxyAccess, error) {
	trimmedToken := strings.TrimSpace(token)
	if trimmedToken == "" {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}

	tokenHash := secrets.HashToken(trimmedToken)
	authKey, err := txs.loadOrganizationLLMProxyKeyByTokenHash(ctx, tokenHash)
	if err != nil {
		return OrganizationProviderProxyAccess{}, err
	}
	allowedProviderIDs, allowedModels, err := txs.listActiveOrganizationProxyAccessScope(ctx, authKey.OrganizationID)
	if err != nil {
		return OrganizationProviderProxyAccess{}, err
	}
	now := time.Now().UTC()
	if _, err := txs.db.Exec(ctx, `
UPDATE organization_llm_proxy_keys
SET last_used_at = $3, updated_at = $3
WHERE organization_id = $1
  AND id = $2
`, authKey.OrganizationID, authKey.ID, now); err != nil {
		return OrganizationProviderProxyAccess{}, fmt.Errorf("update organization llm proxy key last_used_at: %w", err)
	}
	return OrganizationProviderProxyAccess{
		KeyID:              authKey.ID,
		OrganizationID:     authKey.OrganizationID,
		AllowedProviderIDs: allowedProviderIDs,
		AllowedModels:      allowedModels,
		Source:             "persistent_organization",
	}, nil
}

func (txs TxStore) AuthenticateSignedWorkspaceLLMProxyKey(_ context.Context, token string) (OrganizationProviderProxyAccess, error) {
	return txs.authenticateSignedOrganizationProxyKey(token)
}

func (txs TxStore) authenticateSignedOrganizationProxyKey(token string) (OrganizationProviderProxyAccess, error) {
	if !txs.hasOrganizationProxyKeySigningKey() {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}
	raw := strings.TrimPrefix(strings.TrimSpace(token), "lpi_")
	parts := strings.Split(raw, ".")
	if len(parts) != 2 {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}
	encodedPayload := parts[0]
	expectedMAC := hmac.New(sha256.New, txs.feishuCredentialEncryptionKey)
	_, _ = expectedMAC.Write([]byte(encodedPayload))
	if !hmac.Equal([]byte(parts[1]), []byte(hex.EncodeToString(expectedMAC.Sum(nil)))) {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}
	payloadBytes, err := base64.RawURLEncoding.DecodeString(encodedPayload)
	if err != nil {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}
	var payload model.SignedProviderProxyKey
	if err := model.UnmarshalJSONObject("signed provider proxy key", payloadBytes, &payload); err != nil {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}
	if payload.WorkspaceID == "" {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}
	if payload.ExpiresAtUnix > 0 && time.Now().UTC().Unix() > payload.ExpiresAtUnix {
		return OrganizationProviderProxyAccess{}, ErrNotFound
	}

	visibleProviders, err := txs.ListVisibleOrganizationProvidersForWorkspace(context.Background(), payload.WorkspaceID)
	if err != nil {
		return OrganizationProviderProxyAccess{}, err
	}

	allowedProviderIDs := make([]string, 0, len(visibleProviders))
	providerSeen := make(map[string]struct{}, len(visibleProviders))
	modelSeen := make(map[string]struct{})
	allowedModels := make([]string, 0)
	organizationID := ""
	for _, provider := range visibleProviders {
		trimmedProviderID := strings.TrimSpace(provider.ID)
		if trimmedProviderID != "" {
			if _, ok := providerSeen[trimmedProviderID]; !ok {
				providerSeen[trimmedProviderID] = struct{}{}
				allowedProviderIDs = append(allowedProviderIDs, trimmedProviderID)
			}
		}
		if organizationID == "" {
			organizationID = strings.TrimSpace(provider.OrganizationID)
		}
		for _, model := range provider.Models {
			trimmedModelID := strings.TrimSpace(model.ID)
			if trimmedModelID == "" {
				continue
			}
			if _, ok := modelSeen[trimmedModelID]; ok {
				continue
			}
			modelSeen[trimmedModelID] = struct{}{}
			allowedModels = append(allowedModels, trimmedModelID)
		}
	}

	return OrganizationProviderProxyAccess{
		WorkspaceID:        payload.WorkspaceID,
		OrganizationID:     organizationID,
		AllowedProviderIDs: allowedProviderIDs,
		AllowedModels:      allowedModels,
		Source:             "signed_workspace",
	}, nil
}