package server

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

	"go.uber.org/zap"

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

// maxExtractConfigPromptBytes caps the tenant custom extract prompt.
const maxExtractConfigPromptBytes = 8 << 10

const maskedAPIKeySuffix = "********"

// extractConfigResponse is the admin-facing view of one tenant extract
// config. api_key is always masked; source is "custom" (tenant row),
// "default" (process default), or "none" (no config at all).
type extractConfigResponse struct {
	Enabled   bool   `json:"enabled"`
	APIBase   string `json:"api_base,omitempty"`
	APIKey    string `json:"api_key,omitempty"` // masked
	Model     string `json:"model,omitempty"`
	Prompt    string `json:"prompt,omitempty"`
	Protocol  string `json:"protocol,omitempty"`
	Source    string `json:"source"` // "custom" | "default" | "none"
	UpdatedAt string `json:"updated_at,omitempty"`
}

// maskAPIKey renders an API key safe for responses without exposing its exact
// length: keys longer than 4 characters keep the first 4; all masks use the
// same fixed suffix.
func maskAPIKey(k string) string {
	runes := []rune(k)
	if len(runes) <= 4 {
		return maskedAPIKeySuffix
	}
	return string(runes[:4]) + maskedAPIKeySuffix
}

// extractConfigPutRequest carries a decoded PUT body. Nil pointer = field
// omitted; explicit JSON null is rejected at decode time.
type extractConfigPutRequest struct {
	enabled  *bool
	apiBase  *string
	apiKey   *string
	model    *string
	prompt   *string
	protocol *string
	// validatedProtocol binds an out-of-transaction provider preflight to the
	// effective protocol observed after the config row is locked.
	validatedProtocol *string
}

var errExtractConfigProtocolChanged = errors.New("extract config protocol changed during provider validation")

func (r *extractConfigPutRequest) protocolValue() string {
	if r == nil || r.protocol == nil {
		return backend.MediaProviderProtocolOpenAI
	}
	return backend.NormalizeMediaProviderProtocol(*r.protocol)
}

// extractConfigRequestError marks state-dependent validation failures that
// map to 400 (vs backend/tx failures that map to 5xx).
type extractConfigRequestError struct{ msg string }

func (e *extractConfigRequestError) Error() string { return e.msg }

func decodeExtractConfigPutRequest(w http.ResponseWriter, r *http.Request) (*extractConfigPutRequest, error) {
	if r.Body == nil {
		return nil, fmt.Errorf("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 extra struct{}
	if err := dec.Decode(&extra); !errors.Is(err, io.EOF) {
		return nil, fmt.Errorf("invalid JSON body: trailing data")
	}
	if raw == nil {
		return nil, fmt.Errorf("request body must be a JSON object")
	}
	req := &extractConfigPutRequest{}
	for key, value := range raw {
		if len(bytes.TrimSpace(value)) == 0 || bytes.Equal(bytes.TrimSpace(value), []byte("null")) {
			return nil, fmt.Errorf("field %q must not be null", key)
		}
		switch key {
		case "enabled":
			var v bool
			if err := json.Unmarshal(value, &v); err != nil {
				return nil, fmt.Errorf("field \"enabled\" must be a boolean")
			}
			req.enabled = &v
		case "api_base":
			v, err := decodeExtractConfigStringField(key, value)
			if err != nil {
				return nil, err
			}
			req.apiBase = v
		case "api_key":
			v, err := decodeExtractConfigStringField(key, value)
			if err != nil {
				return nil, err
			}
			req.apiKey = v
		case "model":
			v, err := decodeExtractConfigStringField(key, value)
			if err != nil {
				return nil, err
			}
			req.model = v
		case "prompt":
			v, err := decodeExtractConfigStringField(key, value)
			if err != nil {
				return nil, err
			}
			req.prompt = v
		case "protocol":
			v, err := decodeExtractConfigStringField(key, value)
			if err != nil {
				return nil, err
			}
			req.protocol = v
		default:
			return nil, fmt.Errorf("unknown field %q", key)
		}
	}
	return req, nil
}

func decodeExtractConfigStringField(key string, value json.RawMessage) (*string, error) {
	var v string
	if err := json.Unmarshal(value, &v); err != nil {
		return nil, fmt.Errorf("field %q must be a string", key)
	}
	return &v, nil
}

// staticValidate checks request rules that do not depend on the existing row
// state. State-dependent rules (trio completeness, empty strings on enabled
// rows) are enforced inside the write transaction.
func (r *extractConfigPutRequest) staticValidate() error {
	if r.protocol != nil {
		protocol := backend.NormalizeMediaProviderProtocol(*r.protocol)
		if protocol != backend.MediaProviderProtocolOpenAI && protocol != backend.MediaProviderProtocolQwenASR {
			return fmt.Errorf("unsupported extract provider protocol %q", protocol)
		}
		*r.protocol = protocol
	}
	if r.apiKey != nil && strings.Contains(*r.apiKey, "*") {
		return errors.New(`api_key must not contain '*'`)
	}
	if r.prompt != nil && len(*r.prompt) > maxExtractConfigPromptBytes {
		return fmt.Errorf("prompt must be at most %d bytes", maxExtractConfigPromptBytes)
	}
	if r.apiBase != nil && *r.apiBase != "" {
		if err := backend.ValidateTenantAPIBase(*r.apiBase); err != nil {
			return err
		}
	}
	if r.enabled == nil || *r.enabled {
		providerFields := 0
		for _, provided := range []bool{r.apiBase != nil, r.apiKey != nil, r.model != nil} {
			if provided {
				providerFields++
			}
		}
		if providerFields > 0 && providerFields < 3 {
			return errors.New("api_base, api_key and model must be provided together")
		}
		if providerFields == 3 {
			if *r.apiBase == "" {
				return errors.New("api_base must not be empty when enabled")
			}
			if *r.apiKey == "" {
				return errors.New("api_key must not be empty when enabled")
			}
			if *r.model == "" {
				return errors.New("model must not be empty when enabled")
			}
		}
	}
	return nil
}

func validateExtractConfigProtocol(mediaType string, req *extractConfigPutRequest) error {
	if req == nil {
		return nil
	}
	providerTrio := req.apiBase != nil && req.apiKey != nil && req.model != nil
	if mediaType == meta.ExtractMediaTypeAudio {
		if req.protocol != nil && !providerTrio {
			return errors.New("api_base, api_key and model must be provided when changing audio protocol")
		}
	}
	if req.protocol == nil || strings.TrimSpace(*req.protocol) == "" {
		return nil
	}
	protocol := strings.ToLower(strings.TrimSpace(*req.protocol))
	if mediaType != meta.ExtractMediaTypeAudio && protocol != backend.MediaProviderProtocolOpenAI {
		return fmt.Errorf("protocol %q is only supported for audio extract", protocol)
	}
	return nil
}

func (r *extractConfigPutRequest) providerValidationConfig() *backend.OpenAIImageTextExtractorConfig {
	if r.enabled != nil && !*r.enabled {
		return nil
	}
	if r.apiBase == nil || r.apiKey == nil || r.model == nil {
		return nil
	}
	return &backend.OpenAIImageTextExtractorConfig{
		BaseURL: *r.apiBase,
		APIKey:  *r.apiKey,
		Model:   *r.model,
		Timeout: defaultTenantImageExtractValidationTimeout,
	}
}

func providerValidationHTTPStatus(category backend.ImageProviderFailureCategory) int {
	switch category {
	case backend.ImageProviderFailureAuthRejected,
		backend.ImageProviderFailureNotFound,
		backend.ImageProviderFailureRequestRejected,
		backend.ImageProviderFailureInvalidResponse:
		return http.StatusBadRequest
	case backend.ImageProviderFailureTimeout:
		return http.StatusGatewayTimeout
	default:
		return http.StatusBadGateway
	}
}

func (s *Server) handleAdminTenantExtractConfigGet(w http.ResponseWriter, r *http.Request, tenantID string, mediaType 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_extract_config_get")
	if err != nil {
		writeAdminTiDBCloudError(w, r.Context(), err, "get tenant extract config")
		return
	}
	if _, ok := s.authorizedAdminTenant(w, r, tenantID, organizationID); !ok {
		return
	}
	row, err := s.meta.GetTenantExtractConfig(r.Context(), tenantID, mediaType)
	if err != nil && !errors.Is(err, meta.ErrNotFound) {
		status := backendErrorStatus(r.Context(), err)
		writeControlPlaneError(w, status, controlPlaneErrorCode(status, err), "extract config lookup failed")
		return
	}
	if err == nil {
		resp, respErr := s.extractConfigResponseForRow(r.Context(), row)
		if respErr != nil {
			logger.Warn(r.Context(), "admin_tenant_extract_config_render_failed", zap.String("tenant_id", tenantID), zap.Error(respErr))
			writeControlPlaneStatusError(w, http.StatusInternalServerError, "extract config render failed")
			return
		}
		writeJSON(w, http.StatusOK, resp)
		return
	}
	view := s.extractDefaultViews[mediaType]
	switch {
	case !view.configured:
		writeJSON(w, http.StatusOK, extractConfigResponse{Enabled: false, Source: "none"})
	case !view.provider:
		writeJSON(w, http.StatusOK, extractConfigResponse{Enabled: true, Source: "default", Protocol: view.protocol})
	default:
		writeJSON(w, http.StatusOK, extractConfigResponse{
			Enabled:  true,
			Source:   "default",
			APIBase:  view.apiBase,
			APIKey:   maskAPIKey(view.apiKey),
			Model:    view.model,
			Prompt:   view.prompt,
			Protocol: view.protocol,
		})
	}
}

func (s *Server) handleAdminTenantExtractConfigPut(w http.ResponseWriter, r *http.Request, tenantID string, mediaType string) {
	req, err := decodeExtractConfigPutRequest(w, r)
	if err != nil {
		writeControlPlaneStatusError(w, http.StatusBadRequest, err.Error())
		return
	}
	if err := req.staticValidate(); err != nil {
		writeControlPlaneStatusError(w, http.StatusBadRequest, err.Error())
		return
	}
	if err := validateExtractConfigProtocol(mediaType, req); 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_extract_config_set")
	if err != nil {
		writeAdminTiDBCloudError(w, r.Context(), err, "update tenant extract config")
		return
	}
	if _, ok := s.authorizedAdminTenant(w, r, tenantID, organizationID); !ok {
		return
	}
	if providerCfg := req.providerValidationConfig(); providerCfg != nil {
		validationProtocol := req.protocolValue()
		if mediaType == meta.ExtractMediaTypeAudio && req.protocol == nil {
			// Preserve an enabled row's protocol for provider-trio-only updates;
			// absent rows continue to use the OpenAI default. Keep this lookup
			// outside the write transaction so provider I/O is never under a lock.
			row, lookupErr := s.meta.GetTenantExtractConfig(r.Context(), tenantID, mediaType)
			if lookupErr != nil && !errors.Is(lookupErr, meta.ErrNotFound) {
				status := backendErrorStatus(r.Context(), lookupErr)
				writeControlPlaneError(w, status, controlPlaneErrorCode(status, lookupErr), "extract config lookup failed")
				return
			}
			if lookupErr == nil && row.Protocol != "" {
				validationProtocol = backend.NormalizeMediaProviderProtocol(row.Protocol)
			}
		}
		start := time.Now()
		var validateErr error
		if mediaType == meta.ExtractMediaTypeAudio {
			validateErr = backend.ValidateAudioProvider(r.Context(), validationProtocol, backend.AudioTextExtractorConfig{BaseURL: providerCfg.BaseURL, APIKey: providerCfg.APIKey, Model: providerCfg.Model, Prompt: providerCfg.Prompt, Timeout: defaultTenantAudioExtractTimeout})
		} else {
			validator := s.imageExtractProviderValidator
			if validator == nil {
				validator = backend.ValidateOpenAIImageProvider
			}
			validateErr = validator(r.Context(), *providerCfg)
		}
		if validateErr != nil {
			category, providerStatus, classified := backend.ImageProviderFailure(validateErr)
			if !classified {
				category = backend.ImageProviderFailureInvalidResponse
			}
			metrics.RecordTenantOperationWithOrg(tenantID, organizationID, mediaType+"_extract", "provider_validate", string(category), time.Since(start))
			logger.Warn(r.Context(), "admin_tenant_extract_provider_validation_failed",
				zap.String("tenant_id", tenantID),
				zap.String("failure_category", string(category)),
				zap.Int("provider_status", providerStatus))
			safeReason := validateErr.Error()
			if !classified {
				safeReason = "provider returned an invalid response"
			}
			status := providerValidationHTTPStatus(category)
			writeControlPlaneError(w, status, controlPlaneErrorCode(status, validateErr), mediaType+" extraction provider validation failed: "+safeReason)
			return
		}
		metrics.RecordTenantOperationWithOrg(tenantID, organizationID, mediaType+"_extract", "provider_validate", "ok", time.Since(start))
		if mediaType == meta.ExtractMediaTypeAudio {
			req.validatedProtocol = &validationProtocol
		}
	}
	var final *meta.TenantExtractConfig
	txErr := s.meta.InTx(r.Context(), func(tx *sql.Tx) error {
		var applyErr error
		final, applyErr = s.applyExtractConfigPut(r.Context(), tx, tenantID, mediaType, req)
		return applyErr
	})
	if txErr != nil {
		if errors.Is(txErr, meta.ErrNotFound) {
			// The tenant vanished or entered deletion between the
			// authorization precheck and the in-transaction tenant lock.
			writeControlPlaneStatusError(w, http.StatusNotFound, "tenant not found")
			return
		}
		var reqErr *extractConfigRequestError
		if errors.As(txErr, &reqErr) {
			writeControlPlaneStatusError(w, http.StatusBadRequest, reqErr.msg)
			return
		}
		if errors.Is(txErr, errExtractConfigProtocolChanged) {
			writeControlPlaneStatusError(w, http.StatusConflict, "extract config changed during provider validation; retry the request")
			return
		}
		logger.Warn(r.Context(), "admin_tenant_extract_config_update_failed", zap.String("tenant_id", tenantID), zap.Error(txErr))
		status := backendErrorStatus(r.Context(), txErr)
		writeControlPlaneError(w, status, controlPlaneErrorCode(status, txErr), "update extract config failed")
		return
	}
	// Drop this pod's cached resolution so the new row (and any replaced
	// plaintext key) takes effect immediately instead of after a TTL window.
	s.extractResolver.Evict(tenantID)
	resp, err := s.extractConfigResponseForRow(r.Context(), final)
	if err != nil {
		logger.Warn(r.Context(), "admin_tenant_extract_config_render_failed", zap.String("tenant_id", tenantID), zap.Error(err))
		writeControlPlaneStatusError(w, http.StatusInternalServerError, "extract config render failed")
		return
	}
	writeJSON(w, http.StatusOK, resp)
}

// applyExtractConfigPut applies the PUT state transition table inside tx.
// The tenant row is locked first (lock order: tenant row, then config row) so
// a terminal delete committing after the handler's authorization precheck
// cannot leave a re-inserted config row behind; the existing config row is
// then locked with SELECT ... FOR UPDATE so concurrent partial PUTs serialize
// and cannot lose each other's omitted fields.
func (s *Server) applyExtractConfigPut(ctx context.Context, tx *sql.Tx, tenantID, mediaType string, req *extractConfigPutRequest) (*meta.TenantExtractConfig, error) {
	if err := meta.LockTenantForConfigUpdateTx(ctx, tx, tenantID); err != nil {
		return nil, err
	}
	row, err := s.meta.GetTenantExtractConfigForUpdateTx(ctx, tx, tenantID, mediaType)
	if err != nil && !errors.Is(err, meta.ErrNotFound) {
		return nil, err
	}
	existingEnabled := err == nil && row.Enabled
	targetEnabled := false
	if req.enabled == nil {
		if !existingEnabled {
			return nil, &extractConfigRequestError{msg: "enabled is required"}
		}
		targetEnabled = true
	} else {
		targetEnabled = *req.enabled
	}
	now := time.Now().UTC()
	if !targetEnabled {
		if err == nil && !row.Enabled {
			// off-marker + enabled:false is a no-op.
			return row, nil
		}
		next := &meta.TenantExtractConfig{
			TenantID: tenantID, MediaType: mediaType,
			Enabled: false, UpdatedAt: now,
		}
		if err == nil {
			next.CreatedAt = row.CreatedAt
		}
		if err := s.meta.UpsertTenantExtractConfigTx(ctx, tx, next); err != nil {
			return nil, err
		}
		return next, nil
	}
	next := &meta.TenantExtractConfig{
		TenantID: tenantID, MediaType: mediaType,
		Enabled: true, UpdatedAt: now,
	}
	if existingEnabled {
		// Omitted fields stay unchanged; provided fields replace.
		next.CreatedAt = row.CreatedAt
		next.APIBase = row.APIBase
		next.Model = row.Model
		next.APIKeyCipher = row.APIKeyCipher
		next.Prompt = row.Prompt
		next.Protocol = row.Protocol
		if req.apiBase != nil {
			next.APIBase = *req.apiBase
		}
		if req.model != nil {
			next.Model = *req.model
		}
		if req.prompt != nil {
			next.Prompt = *req.prompt // "" normalizes to NULL at write
		}
		if req.protocol != nil {
			next.Protocol = *req.protocol
		}
	} else {
		// Absent or off-marker rows need the full trio to enable.
		if req.apiBase == nil || req.apiKey == nil || req.model == nil {
			return nil, &extractConfigRequestError{msg: "api_base, api_key and model are required when enabling"}
		}
		next.APIBase = *req.apiBase
		next.Model = *req.model
		if req.prompt != nil {
			next.Prompt = *req.prompt
		}
		if req.protocol != nil {
			next.Protocol = *req.protocol
		} else if mediaType == meta.ExtractMediaTypeAudio {
			// New audio rows default to OpenAI. Existing enabled rows preserve
			// their protocol when the field is omitted.
			next.Protocol = backend.MediaProviderProtocolOpenAI
		}
	}
	if next.APIBase == "" {
		return nil, &extractConfigRequestError{msg: "api_base must not be empty when enabled"}
	}
	if next.Model == "" {
		return nil, &extractConfigRequestError{msg: "model must not be empty when enabled"}
	}
	if req.validatedProtocol != nil && backend.NormalizeMediaProviderProtocol(next.Protocol) != *req.validatedProtocol {
		return nil, errExtractConfigProtocolChanged
	}
	if err := backend.ValidateTenantAPIBase(next.APIBase); err != nil {
		return nil, &extractConfigRequestError{msg: err.Error()}
	}
	if req.apiKey != nil {
		if *req.apiKey == "" {
			return nil, &extractConfigRequestError{msg: "api_key must not be empty when enabled"}
		}
		cipher, encErr := s.pool.Encrypt(ctx, []byte(*req.apiKey))
		if encErr != nil {
			return nil, fmt.Errorf("encrypt tenant %s extractor api key: %w", tenantID, encErr)
		}
		next.APIKeyCipher = cipher
	}
	if len(next.APIKeyCipher) == 0 {
		return nil, &extractConfigRequestError{msg: "api_key must not be empty when enabled"}
	}
	if err := s.meta.UpsertTenantExtractConfigTx(ctx, tx, next); err != nil {
		return nil, err
	}
	return next, nil
}

// extractConfigResponseForRow renders a persisted row for GET/PUT responses,
// decrypting the api key only to mask it.
func (s *Server) extractConfigResponseForRow(ctx context.Context, row *meta.TenantExtractConfig) (extractConfigResponse, error) {
	resp := extractConfigResponse{Enabled: row.Enabled, Source: "custom"}
	if !row.Enabled {
		return resp, nil
	}
	resp.APIBase = row.APIBase
	resp.Model = row.Model
	resp.Prompt = row.Prompt
	resp.Protocol = row.Protocol
	if resp.Protocol == "" {
		resp.Protocol = "openai"
	}
	if len(row.APIKeyCipher) > 0 {
		key, err := s.pool.Decrypt(ctx, row.APIKeyCipher)
		if err != nil {
			return extractConfigResponse{}, fmt.Errorf("decrypt tenant %s extractor api key: %w", row.TenantID, err)
		}
		resp.APIKey = maskAPIKey(string(key))
	}
	resp.UpdatedAt = row.UpdatedAt.UTC().Format(time.RFC3339)
	return resp, nil
}