package main

import (
	"context"
	"fmt"
	"strconv"

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

type Client struct {
	c     *api.Client
	debug bool
	model string
}

func NewClient(model string, debug bool) (*Client, error) {

	client, err := api.ClientFromEnvironment()
	if err != nil {
		return nil, err
	}

	return &Client{
		c:     client,
		debug: debug,
		model: model,
	}, nil
}

func (c Client) Translate(ctx context.Context, input string, local, remote Language) (string, error) {
	messages := []api.Message{
		{
			Role: "system",
			Content: fmt.Sprintf(
				TRANSLATE_TEMPLATE, remote.String(), local.String()),
		},
		{
			Role:    "user",
			Content: input,
		},
	}

	text := ""

	req := &api.ChatRequest{
		Model:    c.model,
		Messages: messages,
	}

	respFunc := func(resp api.ChatResponse) error {
		if c.debug {
			fmt.Print(resp.Message.Content)
		}
		text += resp.Message.Content
		return nil
	}

	err := c.c.Chat(ctx, req, respFunc)
	if err != nil {
		return "", nil
	}

	return text, nil
}

func (c Client) Compare(ctx context.Context, t1, t2 string) (int, error) {
	promptInput := fmt.Sprintf(COMPARISON_TEMPLATE, t1, t2)
	log.Debugf("input: %s", promptInput)
	messages := []api.Message{
		{
			Role:    "system",
			Content: COMPARISION_TEMPLATE_SYSTEM,
		},
		{
			Role:    "user",
			Content: promptInput,
		},
	}

	text := ""

	req := &api.ChatRequest{
		Model:    c.model,
		Messages: messages,
	}

	respFunc := func(resp api.ChatResponse) error {
		text += resp.Message.Content
		return nil
	}

	err := c.c.Chat(ctx, req, respFunc)
	if err != nil {
		return -1, err
	}

	score, err := strconv.Atoi(text)

	if err != nil {
		log.Errorf("Failed to extract numeric score from model output: %s", err)
	}

	return score, nil
}