package kb

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
	"unicode/utf8"
)

const (
	// EmbeddingDim BGE-M3 标准向量维度
	EmbeddingDim = 1024

	// embedRequestTokenBudget 单批次 Token 上限
	embedRequestTokenBudget = 7000
	// embedMaxBatchItems 单请求文本条数上限
	embedMaxBatchItems = 16
	// embedMaxItemTokens 单条文本截断上限
	embedMaxItemTokens = 6000
	// embedLogBodyMax 错误截断最大字节
	embedLogBodyMax = 2048
)

// Embedder 向量生成客户端
type Embedder struct {
	Endpoint string
	APIKey   string
	Model    string
	HTTP     *http.Client
}

// NewEmbedder 创建 Embedder 实例
func NewEmbedder(endpoint, key, model string) *Embedder {
	return &Embedder{
		Endpoint: normalizeAPIEndpoint(endpoint, "/embeddings"),
		APIKey:   key,
		Model:    model,
		HTTP:     &http.Client{Timeout: 60 * time.Second},
	}
}

// Embed 批量生成文本嵌入向量
func (e *Embedder) Embed(ctx context.Context, texts []string) ([][]float32, error) {
	if e == nil || e.Endpoint == "" {
		return nil, fmt.Errorf("embedder not configured")
	}
	if len(texts) == 0 {
		return nil, nil
	}

	out := make([][]float32, len(texts))
	for start := 0; start < len(texts); {
		end := nextEmbedBatchEnd(texts, start, embedRequestTokenBudget, embedMaxBatchItems)
		batch := texts[start:end]
		vecs, err := e.embedBatch(ctx, batch)
		if err != nil {
			if isContextLengthError(err) && len(batch) > 1 {
				mid := start + (end-start+1)/2
				if mid <= start {
					mid = start + 1
				}
				if mid >= end {
					mid = end - 1
				}
				left, lerr := e.Embed(ctx, texts[start:mid])
				if lerr != nil {
					return nil, lerr
				}
				right, rerr := e.Embed(ctx, texts[mid:end])
				if rerr != nil {
					return nil, rerr
				}
				copy(out[start:mid], left)
				copy(out[mid:end], right)
				start = end
				continue
			}
			est := estimateTokensBatch(batch)
			return nil, fmt.Errorf("%w (batch=%d est_tokens≈%d)", err, len(batch), est)
		}
		copy(out[start:end], vecs)
		start = end
	}
	return out, nil
}

// EmbedOne 单条文本向量化辅助方法
func (e *Embedder) EmbedOne(ctx context.Context, text string) ([]float32, error) {
	vecs, err := e.Embed(ctx, []string{text})
	if err != nil {
		return nil, err
	}
	if len(vecs) == 0 {
		return nil, fmt.Errorf("empty embedding returned")
	}
	return vecs[0], nil
}

func (e *Embedder) embedBatch(ctx context.Context, texts []string) ([][]float32, error) {
	body := map[string]any{"model": e.Model, "input": texts}
	raw, err := e.post(ctx, e.Endpoint, body)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data []struct {
			Embedding []float32 `json:"embedding"`
			Index     int       `json:"index"`
		} `json:"data"`
	}
	if err := json.Unmarshal(raw, &resp); err != nil {
		return nil, fmt.Errorf("decode embeddings: %s", formatHTTPBodyError(0, raw))
	}
	out := make([][]float32, len(texts))
	for _, d := range resp.Data {
		if d.Index >= 0 && d.Index < len(out) {
			out[d.Index] = padOrTrim(d.Embedding, EmbeddingDim)
		}
	}
	for i, v := range out {
		if v == nil && i < len(resp.Data) {
			out[i] = padOrTrim(resp.Data[i].Embedding, EmbeddingDim)
		}
		if out[i] == nil {
			return nil, fmt.Errorf("missing embedding %d", i)
		}
	}
	return out, nil
}

func (e *Embedder) post(ctx context.Context, url string, body any) ([]byte, error) {
	return postJSON(ctx, e.HTTP, url, e.APIKey, body)
}

func postJSON(ctx context.Context, httpc *http.Client, url, key string, body any) ([]byte, error) {
	b, _ := json.Marshal(body)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	if key != "" {
		req.Header.Set("Authorization", "Bearer "+key)
	}
	resp, err := httpc.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if resp.StatusCode >= 300 {
		return nil, fmt.Errorf("http %s: %s", resp.Status, formatHTTPBodyError(resp.StatusCode, raw))
	}
	return raw, nil
}

func formatHTTPBodyError(status int, raw []byte) string {
	msg := extractAPIErrorMessage(raw)
	if msg != "" {
		return msg
	}
	trim := strings.TrimSpace(string(raw))
	if trim == "" {
		if status > 0 {
			return fmt.Sprintf("empty body (status %d)", status)
		}
		return "empty body"
	}
	if !json.Valid(raw) || strings.HasPrefix(trim, "<") {
		preview := truncateUTF8(trim, 200)
		if status > 0 {
			return fmt.Sprintf("non-JSON response (status %d): %s", status, preview)
		}
		return "non-JSON response: " + preview
	}
	return truncateUTF8(trim, embedLogBodyMax)
}

func extractAPIErrorMessage(raw []byte) string {
	var envelope struct {
		Error struct {
			Message string `json:"message"`
			Type    string `json:"type"`
			Code    any    `json:"code"`
		} `json:"error"`
		Message string `json:"message"`
	}
	if err := json.Unmarshal(raw, &envelope); err != nil {
		return ""
	}
	if m := strings.TrimSpace(envelope.Error.Message); m != "" {
		return m
	}
	return strings.TrimSpace(envelope.Message)
}

func isContextLengthError(err error) bool {
	if err == nil {
		return false
	}
	msg := strings.ToLower(err.Error())
	return strings.Contains(msg, "maximum context length") ||
		strings.Contains(msg, "context length") ||
		strings.Contains(msg, "too many tokens") ||
		(strings.Contains(msg, "input tokens") && strings.Contains(msg, "reduce"))
}

func estimateTokens(s string) int {
	s = strings.TrimSpace(s)
	if s == "" {
		return 1
	}
	tokens := 0
	asciiRun := 0
	flushASCII := func() {
		if asciiRun == 0 {
			return
		}
		tokens += (asciiRun + 3) / 4
		asciiRun = 0
	}
	for _, r := range s {
		if r <= 127 {
			asciiRun++
			continue
		}
		flushASCII()
		tokens += 2
	}
	flushASCII()
	if tokens < 1 {
		return 1
	}
	return tokens
}

func estimateTokensBatch(texts []string) int {
	sum := 0
	for _, t := range texts {
		sum += estimateTokens(t)
	}
	return sum
}

func nextEmbedBatchEnd(texts []string, start, budget, maxItems int) int {
	if maxItems <= 0 {
		maxItems = embedMaxBatchItems
	}
	end := start
	used := 0
	for end < len(texts) && end-start < maxItems {
		t := estimateTokens(texts[end])
		if t > embedMaxItemTokens {
			t = embedMaxItemTokens
		}
		if end > start && used+t > budget {
			break
		}
		used += t
		end++
		if used >= budget {
			break
		}
	}
	if end == start {
		end = start + 1
	}
	return end
}

func normalizeAPIEndpoint(endpoint, suffix string) string {
	endpoint = strings.TrimSpace(endpoint)
	if endpoint == "" {
		return ""
	}
	endpoint = strings.TrimRight(endpoint, "/")
	if strings.HasSuffix(endpoint, suffix) {
		return endpoint
	}
	if strings.HasSuffix(endpoint, "/v1") || strings.Contains(endpoint, "/v1/") {
		return endpoint + suffix
	}
	return endpoint + "/v1" + suffix
}

func padOrTrim(v []float32, n int) []float32 {
	if len(v) == n {
		return v
	}
	out := make([]float32, n)
	copy(out, v)
	return out
}

func truncateUTF8(s string, n int) string {
	if n <= 0 || len(s) <= n {
		return s
	}
	for n > 0 && !utf8.RuneStart(s[n]) {
		n--
	}
	if n <= 0 {
		return ""
	}
	return s[:n]
}