package ollama

import (
	"context"
	"net/http"
	"net/url"
	"strings"

	"github.com/ollama/ollama/api"
)

type OllamaClient struct {
	client api.Client
}

func New(ctx context.Context, host *url.URL, http *http.Client) *OllamaClient {
	client := api.NewClient(host, http)
	return &OllamaClient{client: *client}
}

func (o *OllamaClient) EmbedTexts(ctx context.Context, posts []string) ([][]float64, error) {
	vectors := make([][]float64, 0, len(posts))
	for _, t := range posts {
		resp, err := o.client.Embeddings(ctx, &api.EmbeddingRequest{
			Model:  "nomic-embed-text",
			Prompt: t,
		})
		if err != nil {
			return nil, err
		}
		vectors = append(vectors, resp.Embedding)
	}
	return vectors, nil
}

func (o *OllamaClient) SummarizeStyle(ctx context.Context, posts []string) (string, error) {
	prompt := "Summarize writing style, recurring themes, tone, personality traits of this user:\n\n" +
		strings.Join(posts, "\n") + "\n\nStyle summary:"
	var fullResponse strings.Builder
	err := o.client.Generate(ctx, &api.GenerateRequest{
		Model:  "mistral",
		Prompt: prompt,
	}, func(resp api.GenerateResponse) error {
		fullResponse.WriteString(resp.Response)
		return nil
	})
	if err != nil {
		return "", err
	}

	return fullResponse.String(), nil
}