package server

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"strings"
	"time"
	"unicode/utf8"

	"go.uber.org/zap"

	"github.com/mem9-ai/drive9/pkg/backend"
	"github.com/mem9-ai/drive9/pkg/embedding"
	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/metrics"
	"github.com/mem9-ai/drive9/pkg/tenant"
)

type embeddingConfigResponse struct {
	Enabled    bool   `json:"enabled"`
	APIBase    string `json:"api_base,omitempty"`
	APIKey     string `json:"api_key,omitempty"`
	Model      string `json:"model,omitempty"`
	Source     string `json:"source"`
	Generation uint64 `json:"generation,omitempty"`
	UpdatedAt  string `json:"updated_at,omitempty"`
}

type embeddingConfigPutRequest struct {
	Enabled *bool   `json:"enabled"`
	APIBase *string `json:"api_base"`
	APIKey  *string `json:"api_key"`
	Model   *string `json:"model"`
}

type EmbeddingDefaultProvider struct {
	APIBase string
	APIKey  string
	Model   string
}

var errEmbeddingAPIKeyTooLong = errors.New("api_key is too long")

func decodeEmbeddingConfigPutRequest(w http.ResponseWriter, r *http.Request) (*embeddingConfigPutRequest, error) {
	if r.Body == nil {
		return nil, errors.New("request body is required")
	}
	dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxCredentialProvisionBodyBytes))
	var raw map[string]json.RawMessage
	if err := dec.Decode(&raw); err != nil {
		return nil, fmt.Errorf("invalid JSON body: %w", err)
	}
	var trailing struct{}
	if err := dec.Decode(&trailing); !errors.Is(err, io.EOF) {
		return nil, fmt.Errorf("invalid JSON body: trailing data")
	}
	if raw == nil {
		return nil, errors.New("request body must be a JSON object")
	}
	var req embeddingConfigPutRequest
	for key, value := range raw {
		trimmed := bytes.TrimSpace(value)
		if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) {
			return nil, fmt.Errorf("field %q must not be null", key)
		}
		switch key {
		case "enabled":
			var enabled bool
			if err := json.Unmarshal(value, &enabled); err != nil {
				return nil, errors.New(`field "enabled" must be a boolean`)
			}
			req.Enabled = &enabled
		case "api_base":
			if err := json.Unmarshal(value, &req.APIBase); err != nil {
				return nil, errors.New(`field "api_base" must be a string`)
			}
		case "api_key":
			if err := json.Unmarshal(value, &req.APIKey); err != nil {
				return nil, errors.New(`field "api_key" must be a string`)
			}
		case "model":
			if err := json.Unmarshal(value, &req.Model); err != nil {
				return nil, errors.New(`field "model" must be a string`)
			}
		default:
			return nil, fmt.Errorf("unknown field %q", key)
		}
	}
	if req.Enabled == nil {
		return nil, errors.New("enabled is required")
	}
	if *req.Enabled {
		if req.APIBase == nil || req.APIKey == nil || req.Model == nil {
			return nil, errors.New("api_base, api_key and model are required when enabling")
		}
		*req.APIBase = strings.TrimSpace(*req.APIBase)
		*req.APIKey = strings.TrimSpace(*req.APIKey)
		*req.Model = strings.TrimSpace(*req.Model)
		if *req.APIBase == "" || *req.APIKey == "" || *req.Model == "" {
			return nil, errors.New("api_base, api_key and model must not be empty when enabling")
		}
		if utf8.RuneCountInString(*req.APIBase) > meta.TenantEmbeddingAPIBaseMaxChars {
			return nil, fmt.Errorf("api_base must be at most %d characters", meta.TenantEmbeddingAPIBaseMaxChars)
		}
		if utf8.RuneCountInString(*req.Model) > meta.TenantEmbeddingModelMaxChars {
			return nil, fmt.Errorf("model must be at most %d characters", meta.TenantEmbeddingModelMaxChars)
		}
		if strings.Contains(*req.APIKey, "*") {
			return nil, errors.New("api_key must not contain '*'")
		}
		if err := backend.ValidateTenantAPIBase(*req.APIBase); err != nil {
			return nil, err
		}
		return &req, nil
	}
	if req.APIBase != nil || req.APIKey != nil || req.Model != nil {
		return nil, errors.New("api_base, api_key and model must be omitted when disabling")
	}
	return &req, nil
}

func validateEmbeddingProvider(ctx context.Context, req *embeddingConfigPutRequest) error {
	client, err := embedding.NewOpenAIClient(embedding.OpenAIClientConfig{
		BaseURL: *req.APIBase, APIKey: *req.APIKey, Model: *req.Model,
		Dimensions: embedding.Dimensions, Timeout: defaultEmbeddingTimeout,
		Client: backend.NewTenantExtractorHTTPClient(defaultEmbeddingTimeout),
	})
	if err != nil {
		return err
	}
	_, err = client.EmbedText(ctx, "Drive9 embedding provider validation")
	return err
}

func safeEmbeddingProviderValidationError(err error) string {
	if _, _, safeReason, ok := embedding.ProviderFailure(err); ok {
		return safeReason
	}
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		return "provider request timed out"
	case err == nil:
		return "provider validation failed"
	default:
		return "provider request failed"
	}
}

func embeddingProviderValidationHTTPStatus(err error) int {
	category, _, _, ok := embedding.ProviderFailure(err)
	if ok {
		switch category {
		case embedding.ProviderFailureAuthRejected,
			embedding.ProviderFailureNotFound,
			embedding.ProviderFailureRequestRejected,
			embedding.ProviderFailureInvalidResponse:
			return http.StatusBadRequest
		case embedding.ProviderFailureRateLimited:
			return http.StatusTooManyRequests
		}
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return http.StatusGatewayTimeout
	}
	return http.StatusBadGateway
}

func (s *Server) tenantAppEmbeddingMode(ctx context.Context, t *meta.Tenant) (string, error) {
	if t == nil {
		return "", meta.ErrNotFound
	}
	if tenant.IsSharedSchemaProvider(t.Provider) {
		return meta.TenantEmbeddingModeFTSOnly, nil
	}
	if !tenant.UsesTiDBAutoEmbedding(t.Provider) {
		return "", fmt.Errorf("embedding config is not supported for provider %q", t.Provider)
	}
	persisted := ""
	profile, err := s.meta.GetTenantAutoEmbeddingProfile(ctx, t.ID)
	if err == nil {
		persisted = profile.EmbeddingMode
	} else if !errors.Is(err, meta.ErrNotFound) {
		return "", err
	}
	mode, _, err := meta.ResolveTenantEmbeddingMode(persisted, s.disableDBAutoEmbed)
	return mode, err
}

func (s *Server) handleAdminTenantEmbeddingConfigGet(w http.ResponseWriter, r *http.Request, tenantID string) {
	cred, err := controlPlaneCredentials(r, "", "")
	if err != nil {
		writeControlPlaneError(w, http.StatusBadRequest, controlPlaneErrorCode(http.StatusBadRequest, err), err.Error())
		return
	}
	organizationID, err := s.authorizeTiDBCloudAdminOrganization(r.Context(), cred, "admin_tenant_embedding_config_get")
	if err != nil {
		writeAdminTiDBCloudError(w, r.Context(), err, "get tenant embedding config")
		return
	}
	t, ok := s.authorizedAdminTenant(w, r, tenantID, organizationID)
	if !ok {
		return
	}
	mode, err := s.tenantAppEmbeddingMode(r.Context(), t)
	if err != nil {
		status := backendErrorStatus(r.Context(), err)
		writeControlPlaneError(w, status, controlPlaneErrorCode(status, err), "embedding mode lookup failed")
		return
	}
	if mode == meta.TenantEmbeddingModeAuto {
		writeJSON(w, http.StatusOK, embeddingConfigResponse{Source: "database_auto"})
		return
	}
	row, err := s.meta.GetTenantEmbeddingConfig(r.Context(), tenantID)
	if err == nil {
		resp, renderErr := s.embeddingConfigResponseForRow(r.Context(), row)
		if renderErr != nil {
			writeControlPlaneStatusError(w, http.StatusInternalServerError, "embedding config render failed")
			return
		}
		writeJSON(w, http.StatusOK, resp)
		return
	}
	if !errors.Is(err, meta.ErrNotFound) {
		status := backendErrorStatus(r.Context(), err)
		writeControlPlaneError(w, status, controlPlaneErrorCode(status, err), "embedding config lookup failed")
		return
	}
	if s.embeddingDefaultView == nil {
		writeJSON(w, http.StatusOK, embeddingConfigResponse{Source: "none"})
		return
	}
	writeJSON(w, http.StatusOK, embeddingConfigResponse{
		Enabled: true, APIBase: s.embeddingDefaultView.APIBase,
		APIKey: maskAPIKey(s.embeddingDefaultView.APIKey), Model: s.embeddingDefaultView.Model,
		Source: "default", Generation: s.embeddingDefaultGeneration,
	})
}

func (s *Server) handleAdminTenantEmbeddingConfigPut(w http.ResponseWriter, r *http.Request, tenantID string) {
	req, err := decodeEmbeddingConfigPutRequest(w, r)
	if err != nil {
		writeControlPlaneStatusError(w, http.StatusBadRequest, err.Error())
		return
	}
	cred, err := controlPlaneCredentials(r, "", "")
	if err != nil {
		writeControlPlaneError(w, http.StatusBadRequest, controlPlaneErrorCode(http.StatusBadRequest, err), err.Error())
		return
	}
	organizationID, err := s.authorizeTiDBCloudAdminOrganization(r.Context(), cred, "admin_tenant_embedding_config_set")
	if err != nil {
		writeAdminTiDBCloudError(w, r.Context(), err, "update tenant embedding config")
		return
	}
	t, ok := s.authorizedAdminTenant(w, r, tenantID, organizationID)
	if !ok {
		return
	}
	mode, err := s.tenantAppEmbeddingMode(r.Context(), t)
	if err != nil {
		status := backendErrorStatus(r.Context(), err)
		writeControlPlaneError(w, status, controlPlaneErrorCode(status, err), "embedding mode lookup failed")
		return
	}
	if mode == meta.TenantEmbeddingModeAuto {
		writeControlPlaneStatusError(w, http.StatusConflict, "embedding config is not applicable to database auto-embedding tenants")
		return
	}
	if *req.Enabled {
		started := time.Now()
		validator := s.embeddingProviderValidator
		if validator == nil {
			validator = validateEmbeddingProvider
		}
		if err := validator(r.Context(), req); err != nil {
			category, providerStatus, _, classified := embedding.ProviderFailure(err)
			result := "request_failed"
			if classified {
				result = string(category)
			} else if errors.Is(err, context.DeadlineExceeded) {
				result = "timeout"
			}
			metrics.RecordTenantOperationWithOrg(tenantID, organizationID, "embedding", "provider_validate", result, time.Since(started))
			logger.Warn(r.Context(), "admin_tenant_embedding_provider_validation_failed",
				zap.String("tenant_id", tenantID),
				zap.String("failure_category", result),
				zap.Int("provider_status", providerStatus))
			status := embeddingProviderValidationHTTPStatus(err)
			code := controlPlaneErrorCode(status, err)
			if status == http.StatusTooManyRequests {
				// Upstream provider rate limiting is retryable; keep it
				// machine-distinguishable instead of the generic fallback.
				code = "rate_limited"
			}
			writeControlPlaneError(w, status, code, "embedding provider validation failed: "+safeEmbeddingProviderValidationError(err))
			return
		}
		metrics.RecordTenantOperationWithOrg(tenantID, organizationID, "embedding", "provider_validate", "ok", time.Since(started))
	}
	var final *meta.TenantEmbeddingConfig
	err = s.meta.InTx(r.Context(), func(tx *sql.Tx) error {
		if err := meta.LockTenantForConfigUpdateTx(r.Context(), tx, tenantID); err != nil {
			return err
		}
		existing, err := s.meta.GetTenantEmbeddingConfigForUpdateTx(r.Context(), tx, tenantID)
		if err != nil && !errors.Is(err, meta.ErrNotFound) {
			return err
		}
		final, err = s.buildEmbeddingConfigReplacement(r.Context(), tenantID, req, existing)
		if err != nil {
			return err
		}
		if err := s.meta.UpsertTenantEmbeddingConfigTx(r.Context(), tx, final); err != nil {
			return err
		}
		return meta.InsertTenantNotifyTx(r.Context(), tx, tenantID, meta.TenantNotifyWorkEmbeddingConfigCleanup)
	})
	if err != nil {
		if errors.Is(err, meta.ErrNotFound) {
			writeControlPlaneStatusError(w, http.StatusNotFound, "tenant not found")
			return
		}
		if errors.Is(err, errEmbeddingAPIKeyTooLong) {
			writeControlPlaneStatusError(w, http.StatusBadRequest, errEmbeddingAPIKeyTooLong.Error())
			return
		}
		logger.Warn(r.Context(), "admin_tenant_embedding_config_update_failed", zap.String("tenant_id", tenantID), zap.Error(err))
		status := backendErrorStatus(r.Context(), err)
		writeControlPlaneError(w, status, controlPlaneErrorCode(status, err), "update embedding config failed")
		return
	}
	s.invalidateLocalEmbeddingConfig(tenantID)
	resp, err := s.embeddingConfigResponseForRow(r.Context(), final)
	if err != nil {
		writeControlPlaneStatusError(w, http.StatusInternalServerError, "embedding config render failed")
		return
	}
	writeJSON(w, http.StatusOK, resp)
}

func (s *Server) buildEmbeddingConfigReplacement(ctx context.Context, tenantID string, req *embeddingConfigPutRequest, existing *meta.TenantEmbeddingConfig) (*meta.TenantEmbeddingConfig, error) {
	next := &meta.TenantEmbeddingConfig{TenantID: tenantID, Enabled: *req.Enabled, Generation: 1, ConfigRevision: 1, UpdatedAt: time.Now().UTC()}
	if existing != nil {
		next.CreatedAt = existing.CreatedAt
		next.Generation = existing.Generation
		next.ConfigRevision = existing.ConfigRevision
	}
	if !*req.Enabled {
		if existing != nil && existing.Enabled {
			if next.Generation == math.MaxInt64 || next.ConfigRevision == math.MaxInt64 {
				return nil, fmt.Errorf("embedding config counter exhausted")
			}
			next.Generation++
			next.ConfigRevision++
		} else if existing != nil {
			next.UpdatedAt = existing.UpdatedAt
		}
		return next, nil
	}
	next.APIBase, next.Model = *req.APIBase, *req.Model
	cipher, err := s.pool.Encrypt(ctx, []byte(*req.APIKey))
	if err != nil {
		return nil, fmt.Errorf("encrypt tenant embedding api key: %w", err)
	}
	next.APIKeyCipher = cipher
	if len(next.APIKeyCipher) > meta.TenantEmbeddingAPIKeyCipherMaxBytes {
		return nil, errEmbeddingAPIKeyTooLong
	}
	if existing == nil {
		return next, nil
	}
	oldKey := ""
	if len(existing.APIKeyCipher) > 0 {
		plain, err := s.pool.Decrypt(ctx, existing.APIKeyCipher)
		if err != nil {
			return nil, fmt.Errorf("decrypt existing tenant embedding api key: %w", err)
		}
		oldKey = string(plain)
	}
	contractChanged := !existing.Enabled || existing.APIBase != next.APIBase || existing.Model != next.Model
	configChanged := contractChanged || oldKey != *req.APIKey
	if contractChanged && next.Generation == math.MaxInt64 {
		return nil, fmt.Errorf("embedding generation exhausted")
	}
	if configChanged && next.ConfigRevision == math.MaxInt64 {
		return nil, fmt.Errorf("embedding config revision exhausted")
	}
	if contractChanged {
		next.Generation++
	}
	if configChanged {
		next.ConfigRevision++
	} else {
		next.UpdatedAt = existing.UpdatedAt
	}
	return next, nil
}

func (s *Server) embeddingConfigResponseForRow(ctx context.Context, row *meta.TenantEmbeddingConfig) (embeddingConfigResponse, error) {
	resp := embeddingConfigResponse{Enabled: row.Enabled, Source: "custom", Generation: row.Generation}
	if !row.Enabled {
		return resp, nil
	}
	plain, err := s.pool.Decrypt(ctx, row.APIKeyCipher)
	if err != nil {
		return embeddingConfigResponse{}, err
	}
	resp.APIBase, resp.Model, resp.APIKey = row.APIBase, row.Model, maskAPIKey(string(plain))
	resp.UpdatedAt = row.UpdatedAt.UTC().Format(time.RFC3339)
	return resp, nil
}

func (s *Server) invalidateLocalEmbeddingConfig(tenantID string) {
	if s.embeddingResolver != nil {
		s.embeddingResolver.Evict(tenantID)
	}
	if s.pool != nil {
		s.pool.Invalidate(tenantID)
	}
}