package gollama

import (
	"context"

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

// CompleteStream streams completion tokens through fn.
func (c *Client) CompleteStream(ctx context.Context, prompt string, fn func(Chunk) error, opts ...CallOption) error {
	cfg, err := c.mergeCallConfig(opts...)
	if err != nil {
		return err
	}
	optMap, err := cfg.optionsMap()
	if err != nil {
		return wrapError("complete_stream", cfg.model, err)
	}

	stream := true
	req := &api.GenerateRequest{
		Model:     cfg.model,
		Prompt:    prompt,
		System:    cfg.system,
		Stream:    &stream,
		Options:   optMap,
		KeepAlive: cfg.keepAlive,
		Think:     cfg.think,
		Format:    cfg.format,
		Truncate:  cfg.truncate,
		Shift:     cfg.shift,
	}

	return wrapError("complete_stream", cfg.model, c.api.Generate(ctx, req, func(resp api.GenerateResponse) error {
		return fn(Chunk{
			Text:     resp.Response,
			Thinking: resp.Thinking,
			Done:     resp.Done,
			Metrics:  resp.Metrics,
		})
	}))
}

// ChatStream streams chat response tokens through fn.
func (c *Client) ChatStream(ctx context.Context, messages []Message, fn func(Chunk) error, opts ...CallOption) error {
	cfg, err := c.mergeCallConfig(opts...)
	if err != nil {
		return err
	}
	optMap, err := cfg.optionsMap()
	if err != nil {
		return wrapError("chat_stream", cfg.model, err)
	}

	msgs := toAPIMessages(messages)
	if cfg.system != "" {
		msgs = append([]api.Message{{Role: "system", Content: cfg.system}}, msgs...)
	}

	stream := true
	req := &api.ChatRequest{
		Model:     cfg.model,
		Messages:  msgs,
		Stream:    &stream,
		Options:   optMap,
		KeepAlive: cfg.keepAlive,
		Think:     cfg.think,
		Format:    cfg.format,
		Truncate:  cfg.truncate,
		Shift:     cfg.shift,
	}

	return wrapError("chat_stream", cfg.model, c.api.Chat(ctx, req, func(resp api.ChatResponse) error {
		return fn(Chunk{
			Text:     resp.Message.Content,
			Thinking: resp.Message.Thinking,
			Done:     resp.Done,
			Metrics:  resp.Metrics,
		})
	}))
}