// Package llm is a minimal client for OpenAI-compatible chat APIs, as served
// by llama.cpp, oMLX, Ollama, LM Studio, Lemonade, vLLM and friends.
package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// Message is one chat turn.
type Message struct {
	Role    string `json:"role"` // "system", "user" or "assistant"
	Content string `json:"content"`
}

// Client talks to a single OpenAI-compatible endpoint.
type Client struct {
	BaseURL string // e.g. http://127.0.0.1:8000/v1
	APIKey  string
	Model   string // empty: use the first model the server lists
	HTTP    *http.Client
}

func (c *Client) httpClient() *http.Client {
	if c.HTTP != nil {
		return c.HTTP
	}
	return http.DefaultClient
}

func (c *Client) newRequest(ctx context.Context, method, path string, body io.Reader) (*http.Request, error) {
	req, err := http.NewRequestWithContext(ctx, method, c.BaseURL+path, body)
	if err != nil {
		return nil, err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.APIKey)
	}
	return req, nil
}

// Models lists the model ids the server offers. It doubles as a health check.
func (c *Client) Models(ctx context.Context) ([]string, error) {
	req, err := c.newRequest(ctx, http.MethodGet, "/models", nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.httpClient().Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if err := checkStatus(resp); err != nil {
		return nil, err
	}
	var out struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return nil, fmt.Errorf("decode models: %w", err)
	}
	ids := make([]string, len(out.Data))
	for i, m := range out.Data {
		ids[i] = m.ID
	}
	return ids, nil
}

// ResolveModel returns the configured model, or the first one the server
// lists when none is configured.
func (c *Client) ResolveModel(ctx context.Context) (string, error) {
	if c.Model != "" {
		return c.Model, nil
	}
	ids, err := c.Models(ctx)
	if err != nil {
		return "", err
	}
	if len(ids) == 0 {
		return "", errors.New("model server lists no models; set LLM_MODEL")
	}
	return ids[0], nil
}

// Stream sends the conversation and calls onDelta for every chunk of the
// reply as it arrives. It returns when the reply is complete, the context is
// cancelled, or onDelta returns an error.
func (c *Client) Stream(ctx context.Context, msgs []Message, onDelta func(string) error) error {
	model, err := c.ResolveModel(ctx)
	if err != nil {
		return err
	}
	body, err := json.Marshal(map[string]any{
		"model":    model,
		"messages": msgs,
		"stream":   true,
	})
	if err != nil {
		return err
	}
	req, err := c.newRequest(ctx, http.MethodPost, "/chat/completions", bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "text/event-stream")
	resp, err := c.httpClient().Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := checkStatus(resp); err != nil {
		return err
	}

	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 64*1024), 1024*1024)
	for sc.Scan() {
		data, ok := strings.CutPrefix(sc.Text(), "data:")
		if !ok {
			continue // blank separators, comments, "event:" lines
		}
		data = strings.TrimSpace(data)
		if data == "[DONE]" {
			return nil
		}
		var chunk struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
			Error *struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			return fmt.Errorf("decode stream chunk: %w", err)
		}
		if chunk.Error != nil {
			return fmt.Errorf("model server: %s", chunk.Error.Message)
		}
		for _, ch := range chunk.Choices {
			if ch.Delta.Content == "" {
				continue
			}
			if err := onDelta(ch.Delta.Content); err != nil {
				return err
			}
		}
	}
	if err := sc.Err(); err != nil {
		return err
	}
	return nil // stream ended without [DONE]; treat what we got as the reply
}

func checkStatus(resp *http.Response) error {
	if resp.StatusCode < 300 {
		return nil
	}
	msg, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
	return fmt.Errorf("model server returned %s: %s", resp.Status, bytes.TrimSpace(msg))
}