package service

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/caarlos0/httperr"
	"github.com/go-resty/resty/v2"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	"github.com/tidbcloud/inference-service/internal/config"
	"github.com/tidbcloud/inference-service/internal/metrics"
)

// CohereRequest represents a request to Cohere models
type CohereRequest map[string]any

// CohereResponse represents a response from Cohere models
type CohereResponse struct {
	Embeddings [][]float32 `json:"embeddings"`
	Usage      struct {
		InputTokens int `json:"input_tokens"`
	} `json:"usage"`
}

// TitanRequest represents a request to Titan models
type TitanRequest map[string]any

// TitanResponse represents a response from Titan models
type TitanResponse struct {
	Embedding []float32 `json:"embedding"`
}

// BedrockService handles AWS Bedrock integration via HTTP calls
type BedrockService struct {
	client   *resty.Client
	apiToken string
	region   string
}

func escapeBedrockPathLabel(value string) string {
	const hex = "0123456789ABCDEF"
	var b strings.Builder
	b.Grow(len(value))
	for i := 0; i < len(value); i++ {
		c := value[i]
		if (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-' || c == '.' || c == '_' || c == '~' {
			b.WriteByte(c)
			continue
		}
		b.WriteByte('%')
		b.WriteByte(hex[c>>4])
		b.WriteByte(hex[c&15])
	}
	return b.String()
}

func (s *BedrockService) invokeModelURL(modelID string) string {
	return fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com/model/%s/invoke",
		s.region, escapeBedrockPathLabel(modelID))
}

func applyAllowedOpts(request map[string]any, params *EmbeddingRequest, modelConfig *config.ModelConfig) error {
	if params.OtherOpts == nil {
		return nil
	}
	for key, value := range params.OtherOpts {
		allowed := false
		for _, allowedOpt := range modelConfig.AllowedOpts {
			if key == allowedOpt {
				allowed = true
				break
			}
		}
		if !allowed {
			return httperr.Errorf(http.StatusBadRequest, "Option key '%s' is not allowed for model '%s'", key, params.Model)
		}
		request[key] = value
	}
	return nil
}

// NewBedrockService creates a new BedrockService
func NewBedrockService(apiToken, region string) *BedrockService {
	client := resty.New()
	if apiToken == "" && !testing.Testing() {
		zap.L().Fatal("AWS Bedrock API token is required, configure it in the environment variable AWS_BEARER_TOKEN_BEDROCK or in .env file")
	}
	return &BedrockService{
		client:   client,
		apiToken: apiToken,
		region:   region,
	}
}

// GenerateEmbeddings generates embeddings for the given request
func (s *BedrockService) GenerateEmbeddings(ctx context.Context, params *EmbeddingRequest) ([][]float32, int64, error) {
	modelConfig, exists := config.Global().Models[params.Model]
	if !exists {
		// should not happen
		return nil, 0, fmt.Errorf("unexpected model '%s'", params.Model)
	}
	if modelConfig.IsTitanModel() {
		return s.generateTitanEmbeddings(ctx, params, &modelConfig)
	} else if modelConfig.IsCohereModel() {
		return s.generateCohereEmbeddings(ctx, params, &modelConfig)
	}
	// should not happen
	return nil, 0, fmt.Errorf("unable to find model provider for '%s'", params.Model)
}

// generateTitanEmbeddings handles Titan model embedding generation with batching
func (s *BedrockService) generateTitanEmbeddings(ctx context.Context, params *EmbeddingRequest, modelConfig *config.ModelConfig) ([][]float32, int64, error) {
	// Titan models don't support batching, so we need to make individual requests
	embeddings := make([][]float32, len(params.Texts))
	tokens := make([]int64, len(params.Texts))
	errors := make([]error, len(params.Texts))

	g, ctx := errgroup.WithContext(ctx)

	// Process each text concurrently
	for i, text := range params.Texts {
		i, text := i, text // capture loop variables
		g.Go(func() error {
			embedding, tokenCount, err := s.generateSingleTitanEmbedding(ctx, text, params, modelConfig)
			if err != nil {
				errors[i] = err
				// Still record token usage for successful requests
				return err
			}
			embeddings[i] = embedding
			tokens[i] = tokenCount
			return nil
		})
	}

	// Wait for all goroutines to complete
	err := g.Wait()

	// Calculate total tokens from successful requests
	var totalTokens int64
	for _, tokenCount := range tokens {
		totalTokens += tokenCount
	}
	costMicroDollars := totalTokens * int64(modelConfig.MicroDollarPerKToken) / 1000

	// Return error if any occurred, but still return the token usage
	if err != nil {
		return nil, costMicroDollars, err
	}

	return embeddings, costMicroDollars, nil
}

// generateSingleTitanEmbedding generates a single embedding for Titan models
func (s *BedrockService) generateSingleTitanEmbedding(ctx context.Context, text string, params *EmbeddingRequest, modelConfig *config.ModelConfig) ([]float32, int64, error) {
	// Prepare Titan request
	request := TitanRequest{
		"inputText": text,
	}

	if err := applyAllowedOpts(map[string]any(request), params, modelConfig); err != nil {
		return nil, 0, err
	}

	invokeModelID := modelConfig.BedrockModelID
	if modelConfig.BedrockInvokeModelID != "" {
		invokeModelID = modelConfig.BedrockInvokeModelID
	}
	url := s.invokeModelURL(invokeModelID)

	// Make HTTP request with provider latency tracking
	providerStart := time.Now()
	var response TitanResponse
	resp, err := s.client.R().
		SetContext(ctx).
		SetHeader("Authorization", "Bearer "+s.apiToken).
		SetHeader("Accept", "*/*").
		SetHeader("Content-Type", "application/json").
		SetBody(request).
		SetResult(&response).
		Post(url)

	providerDuration := time.Since(providerStart).Seconds()
	metrics.EmbeddingProviderLatency.WithLabelValues(params.Model).Observe(providerDuration)

	if err != nil {
		zap.L().Error("Failed to call Bedrock API for Titan",
			zap.String("model", params.Model),
			zap.String("modelID", invokeModelID),
			zap.Error(err))
		return nil, 0, err
	}

	if resp.StatusCode() != 200 {
		var errMsg string
		type ErrorResponse struct {
			Message string `json:"message"`
		}
		var errorResponse ErrorResponse
		if err := json.Unmarshal(resp.Body(), &errorResponse); err == nil && errorResponse.Message != "" {
			errMsg = errorResponse.Message
		} else {
			errMsg = string(resp.Body())
		}
		if resp.StatusCode() == http.StatusBadRequest {
			// Likely caused by invalid request parameters from users
			return nil, 0, httperr.Errorf(http.StatusBadRequest, "%s", errMsg)
		}
		zap.L().Error("Bedrock API returned error for Titan",
			zap.String("model", params.Model),
			zap.String("modelID", invokeModelID),
			zap.Int("status_code", resp.StatusCode()),
			zap.String("response", string(resp.Body())))
		return nil, 0, fmt.Errorf("failed to generate embedding: %s", errMsg)
	}

	// Extract token count from headers
	tokenCount := int64(0)
	if tokenHeader := resp.Header().Get("x-amzn-bedrock-input-token-count"); tokenHeader != "" {
		if count, err := strconv.Atoi(tokenHeader); err == nil {
			tokenCount = int64(count)
		}
	}

	// If no token count in header, estimate from text length
	if tokenCount == 0 {
		zap.L().Warn("Failed to extract token count from Bedrock response, estimating from text length",
			zap.String("model", params.Model),
			zap.String("modelID", invokeModelID))
		// More accurate estimation: ~4 chars per token, minimum 1 token per text
		estimatedTokens := len(text) / 4
		if estimatedTokens == 0 {
			estimatedTokens = 1 // Minimum 1 token per text
		}
		tokenCount = int64(estimatedTokens)
	}

	return response.Embedding, tokenCount, nil
}

// generateCohereEmbeddings handles Cohere model embedding generation
func (s *BedrockService) generateCohereEmbeddings(ctx context.Context, params *EmbeddingRequest, modelConfig *config.ModelConfig) ([][]float32, int64, error) {
	// Prepare Cohere request
	request := CohereRequest{
		"texts": params.Texts,
	}

	if err := applyAllowedOpts(map[string]any(request), params, modelConfig); err != nil {
		return nil, 0, err
	}

	invokeModelID := modelConfig.BedrockModelID
	if modelConfig.BedrockInvokeModelID != "" {
		invokeModelID = modelConfig.BedrockInvokeModelID
	}
	url := s.invokeModelURL(invokeModelID)

	// Make HTTP request with provider latency tracking
	providerStart := time.Now()
	var response CohereResponse
	resp, err := s.client.R().
		SetContext(ctx).
		SetHeader("Authorization", "Bearer "+s.apiToken).
		SetHeader("Accept", "*/*").
		SetHeader("Content-Type", "application/json").
		SetBody(request).
		SetResult(&response).
		Post(url)

	providerDuration := time.Since(providerStart).Seconds()
	metrics.EmbeddingProviderLatency.WithLabelValues(params.Model).Observe(providerDuration)

	if err != nil {
		zap.L().Error("Failed to call Bedrock API for Cohere",
			zap.String("model", params.Model),
			zap.String("modelID", invokeModelID),
			zap.Error(err))
		return nil, 0, err
	}

	if resp.StatusCode() != 200 {
		var errMsg string
		type ErrorResponse struct {
			Message string `json:"message"`
		}
		var errorResponse ErrorResponse
		if err := json.Unmarshal(resp.Body(), &errorResponse); err == nil && errorResponse.Message != "" {
			errMsg = errorResponse.Message
		} else {
			errMsg = string(resp.Body())
		}
		if resp.StatusCode() == http.StatusBadRequest {
			// Likely caused by invalid request parameters from users
			return nil, 0, httperr.Errorf(http.StatusBadRequest, "%s", errMsg)
		}
		zap.L().Error("Bedrock API returned error for Cohere",
			zap.String("model", params.Model),
			zap.String("modelID", invokeModelID),
			zap.Int("status_code", resp.StatusCode()),
			zap.String("response", string(resp.Body())))
		return nil, 0, fmt.Errorf("failed to generate embedding: %s", errMsg)
	}

	// Extract token count from headers
	tokenCount := 0
	if tokenHeader := resp.Header().Get("x-amzn-bedrock-input-token-count"); tokenHeader != "" {
		if count, err := strconv.Atoi(tokenHeader); err == nil {
			tokenCount = count
		}
	}

	// If no token count in header, estimate from text length
	if tokenCount == 0 {
		zap.L().Warn("Failed to extract token count from Bedrock response, estimating from text length",
			zap.String("model", params.Model),
			zap.String("modelID", invokeModelID))
		for _, text := range params.Texts {
			// More accurate estimation: ~4 chars per token, minimum 1 token per text
			estimatedTokens := len(text) / 4
			if estimatedTokens == 0 {
				estimatedTokens = 1 // Minimum 1 token per text
			}
			tokenCount += estimatedTokens
		}
	}

	// Calculate cost in micro-dollars
	costMicroDollars := int64(float64(tokenCount) * float64(modelConfig.MicroDollarPerKToken) / 1000)
	return response.Embeddings, costMicroDollars, nil
}