package embedding

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"
)

type OpenAIProvider struct {
	apiKey, model string
	dims          int
	endpoint      string
	client        *http.Client
}

func NewOpenAIProvider(apiKey, model string, dims int) *OpenAIProvider {
	if model == "" {
		model = "text-embedding-3-small"
	}
	if dims <= 0 {
		dims = 1536
	}
	return &OpenAIProvider{apiKey: apiKey, model: model, dims: dims, endpoint: "https://api.openai.com/v1/embeddings", client: &http.Client{Timeout: 30 * time.Second}}
}
func (p *OpenAIProvider) Dimensions() int      { return p.dims }
func (p *OpenAIProvider) Model() string        { return p.model }
func (p *OpenAIProvider) ProviderName() string { return "openai" }
func (p *OpenAIProvider) Embed(ctx context.Context, input []string) ([]Vector, error) {
	if p.apiKey == "" {
		return nil, fmt.Errorf("OPENAI_API_KEY is required")
	}
	body, _ := json.Marshal(map[string]any{"model": p.model, "input": input, "dimensions": p.dims})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+p.apiKey)
	req.Header.Set("Content-Type", "application/json")
	resp, err := p.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("openai embeddings request failed: status %d", resp.StatusCode)
	}
	var parsed struct {
		Data []struct {
			Embedding []float32 `json:"embedding"`
		} `json:"data"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
		return nil, err
	}
	out := make([]Vector, 0, len(parsed.Data))
	for _, item := range parsed.Data {
		if len(item.Embedding) != p.dims {
			return nil, fmt.Errorf("embedding dimension mismatch")
		}
		out = append(out, Vector(item.Embedding))
	}
	return out, nil
}