// Package llm is a minimal client for OpenAI-compatible chat completion APIs,
// with streaming and tool calls.
package llm

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"
	"time"
)

type Client struct {
	BaseURL  string // e.g. "http://localhost:8080/v1"
	APIKey   string // optional for local servers
	Model    string
	Thinking bool // allow reasoning/thinking output on models that support it
	// Sampling settings. Greedy decoding (temperature 0) makes many models,
	// Qwen3 among them, degrade and loop. TopP and TopK are omitted when 0.
	Temperature float64
	TopP        float64
	TopK        int
	HTTP        *http.Client
	Logger      *log.Logger // if non-nil, requests and replies are logged
	// MaxTokens caps each reply, so a model stuck in a repetition loop cannot
	// generate forever. It is sent as max_tokens (which on most servers also
	// counts reasoning), and enforced locally on streamed output in case the
	// server ignores it. 0 means no limit.
	MaxTokens int
	// ContextWindow is how many tokens the model can take in, prompt and
	// reply together; 0 if unknown.
	ContextWindow int

	// What the server accepts besides the standard OpenAI fields, since
	// some (OpenAI's own among them) reject a request with any field they
	// do not know.
	//
	// ThinkingParam is how Thinking is sent: ThinkingTemplate, as
	// enable_thinking both at the top level (SGLang) and in
	// chat_template_kwargs (vLLM, llama.cpp), or "" for not at all.
	ThinkingParam string
	// NoSampling leaves out temperature, top_p and top_k, for servers that
	// take only their own (OpenAI's reasoning models, current Claude models).
	NoSampling bool
	// MaxTokensParam names the field MaxTokens is sent in: "max_tokens", or
	// "max_completion_tokens", which OpenAI's reasoning models require.
	MaxTokensParam string
}

// ThinkingTemplate is the ThinkingParam that sends enable_thinking for the
// chat template.
const ThinkingTemplate = "chat_template_kwargs"

func New(baseURL, apiKey, model string) *Client {
	return &Client{
		BaseURL: strings.TrimSuffix(baseURL, "/"),
		APIKey:  apiKey,
		Model:   model,
		// A redirect would carry the key (and the conversation) somewhere
		// else: Go keeps the Authorization header for the same host name
		// even on another port or over plain http.
		HTTP: &http.Client{Timeout: 10 * time.Minute, CheckRedirect: func(req *http.Request, _ []*http.Request) error {
			return fmt.Errorf("%w to %s; set base_url to where it leads", ErrRedirect, req.URL.Redacted())
		}},
		ThinkingParam:  ThinkingTemplate,
		MaxTokensParam: "max_tokens",
	}
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
	// Reasoning is the thinking that came before an assistant message, sent
	// back so a server whose chat template keeps it can continue from it.
	Reasoning string `json:"reasoning_content,omitempty"`
	// ToolCalls are the calls an assistant message makes; ToolCallID ties a
	// "tool" message's result to its call.
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
}

// ToolCall is one call of a tool the model made.
type ToolCall struct {
	ID       string `json:"id"`
	Type     string `json:"type"` // "function"
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"` // JSON
	} `json:"function"`
}

// Tool describes a function the model may call.
type Tool struct {
	Name        string
	Description string
	Parameters  map[string]any // a JSON schema for the arguments object
}

// toolParam is a tool as sent in a request.
type toolParam struct {
	Type     string `json:"type"`
	Function struct {
		Name        string         `json:"name"`
		Description string         `json:"description,omitempty"`
		Parameters  map[string]any `json:"parameters,omitempty"`
	} `json:"function"`
}

type chatRequest struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Temperature *float64  `json:"temperature,omitempty"`
	TopP        float64   `json:"top_p,omitempty"`
	TopK        int       `json:"top_k,omitempty"` // llama.cpp, vLLM and MLX extension
	// Controls reasoning/thinking output on models that support it (e.g. Qwen3),
	// which slows responses considerably when enabled. SGLang reads the
	// top-level field; vLLM and llama.cpp read it from chat_template_kwargs.
	EnableThinking      *bool          `json:"enable_thinking,omitempty"`
	ChatTemplateKwargs  map[string]any `json:"chat_template_kwargs,omitempty"`
	Stream              bool           `json:"stream,omitempty"`
	StreamOptions       *streamOptions `json:"stream_options,omitempty"`
	MaxTokens           int            `json:"max_tokens,omitempty"`
	MaxCompletionTokens int            `json:"max_completion_tokens,omitempty"`
	Tools               []toolParam    `json:"tools,omitempty"`
	ToolChoice          string         `json:"tool_choice,omitempty"` // "auto" or "none"
}

// streamOptions asks a streaming server to report the tokens used, in a
// last chunk.
type streamOptions struct {
	IncludeUsage bool `json:"include_usage"`
}

// Usage is the tokens a completion took, as the server counts them.
type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
}

// valid reports whether the server sent plausible counts.
func (u *Usage) valid() bool {
	return u != nil && u.PromptTokens > 0 && u.CompletionTokens >= 0
}

// ErrRedirect reports a server's redirect, which the client does not
// follow.
var ErrRedirect = errors.New("the server redirected the request, which blinker does not follow")

// ErrTruncated reports a reply cut off by the MaxTokens limit before the model
// finished. Complete returns the partial reply alongside it.
var ErrTruncated = errors.New("reply hit the length limit")

// bytesPerToken bounds the local fallback limit on streamed output: generous,
// so it only trips when the server ignored max_tokens.
const bytesPerToken = 8

// maxReplyBytes bounds a reply however MaxTokens is set, so a server that
// never stops cannot exhaust memory.
const maxReplyBytes = 64 << 20

type chatResponse struct {
	Choices []struct {
		Message      Message `json:"message"`
		FinishReason string  `json:"finish_reason"`
	} `json:"choices"`
	Usage *Usage `json:"usage"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// Result is one completion: the reply, the reasoning before it, and the
// tool calls it makes.
type Result struct {
	Reply     string
	Reasoning string
	Calls     []ToolCall
	Usage     Usage // zero if the server did not say
}

// Message returns the result as the assistant message to add to the
// conversation.
func (r Result) Message() Message {
	return Message{Role: "assistant", Content: r.Reply, Reasoning: r.Reasoning, ToolCalls: r.Calls}
}

// Complete makes one completion with the tools the model may call, streaming
// the reply to onToken (if non-nil) as it is generated; reasoning marks
// thinking output that is not part of the reply. Transient failures (network
// errors, 408/429/5xx) are retried with backoff, as long as nothing has been
// streamed yet, since a small local server may be busy loading the model or
// serving another request.
func (c *Client) Complete(ctx context.Context, messages []Message, tools []Tool, onToken func(text string, reasoning bool)) (Result, error) {
	var lastErr error
	for attempt := range 3 {
		if attempt > 0 {
			if c.Logger != nil {
				c.Logger.Printf("retrying after error: %v", lastErr)
			}
			select {
			case <-ctx.Done():
				return Result{}, ctx.Err()
			case <-time.After(time.Duration(attempt) * 5 * time.Second):
			}
		}
		streamed := false
		emit := func(text string, reasoning bool) {
			streamed = true
			if onToken != nil {
				onToken(text, reasoning)
			}
		}
		r, retryable, err := c.chatOnce(ctx, messages, tools, emit)
		if err == nil {
			// Some servers leave out a call's id, which its result must
			// name; it need only be unique in the conversation.
			for i := range r.Calls {
				if r.Calls[i].ID == "" {
					r.Calls[i].ID = fmt.Sprintf("call_%d_%d", len(messages), i)
				}
				r.Calls[i].Type = "function"
			}
			return r, nil
		}
		if !retryable || streamed || ctx.Err() != nil {
			return r, err
		}
		lastErr = err
	}
	return Result{}, lastErr
}

// chatOnce makes one streaming request.
func (c *Client) chatOnce(ctx context.Context, messages []Message, tools []Tool, emit func(string, bool)) (r Result, retryable bool, err error) {
	creq := chatRequest{
		Model:         c.Model,
		Messages:      messages,
		Stream:        true,
		StreamOptions: &streamOptions{IncludeUsage: true},
	}
	if !c.NoSampling {
		temp := c.Temperature
		creq.Temperature, creq.TopP, creq.TopK = &temp, c.TopP, c.TopK
	}
	if c.ThinkingParam == ThinkingTemplate {
		thinking := c.Thinking
		creq.EnableThinking = &thinking
		creq.ChatTemplateKwargs = map[string]any{"enable_thinking": thinking}
	}
	if c.MaxTokensParam == "max_completion_tokens" {
		creq.MaxCompletionTokens = c.MaxTokens
	} else {
		creq.MaxTokens = c.MaxTokens
	}
	for _, t := range tools {
		var p toolParam
		p.Type = "function"
		p.Function.Name, p.Function.Description, p.Function.Parameters = t.Name, t.Description, t.Parameters
		creq.Tools = append(creq.Tools, p)
	}
	if len(tools) > 0 {
		creq.ToolChoice = "auto"
	}
	body, err := json.Marshal(creq)
	if err != nil {
		return Result{}, false, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return Result{}, false, err
	}
	req.Header.Set("Content-Type", "application/json")
	if c.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.APIKey)
	}

	if c.Logger != nil {
		dump, err := httputil.DumpRequestOut(req, true)
		if err != nil {
			dump = []byte("dump failed: " + err.Error())
		}
		c.Logger.Printf("request:\n%s", redact(dump, c.APIKey))
	}
	start := time.Now()
	resp, err := c.HTTP.Do(req)
	if err != nil {
		// A network error is worth retrying: the server may be down briefly.
		return Result{}, !errors.Is(err, ErrRedirect), err
	}
	defer resp.Body.Close()

	// A server that honors stream=true answers with server-sent events; some
	// ignore it and send plain JSON, which falls through to the path below.
	if resp.StatusCode == http.StatusOK &&
		strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") {
		return c.readStream(resp, start, emit)
	}

	respBody, err := io.ReadAll(io.LimitReader(resp.Body, 10<<20))
	if err != nil {
		return Result{}, true, err
	}
	if c.Logger != nil {
		head, err := httputil.DumpResponse(resp, false) // body already read; append it below
		if err != nil {
			head = []byte("dump failed: " + err.Error() + "\n")
		}
		c.Logger.Printf("response in %v:\n%s%s", time.Since(start).Round(time.Millisecond), head, respBody)
	}

	// Timeouts, rate limiting, and server errors are worth retrying; other
	// statuses (bad request, auth) will not get better on their own.
	transient := resp.StatusCode == http.StatusRequestTimeout ||
		resp.StatusCode == http.StatusTooManyRequests ||
		resp.StatusCode >= 500

	var parsed chatResponse
	if err := json.Unmarshal(respBody, &parsed); err != nil {
		return Result{}, transient, fmt.Errorf("API returned non-JSON response (status %d): %.200s", resp.StatusCode, respBody)
	}
	if parsed.Error != nil {
		return Result{}, transient, fmt.Errorf("API error (status %d): %s", resp.StatusCode, parsed.Error.Message)
	}
	if resp.StatusCode != http.StatusOK {
		return Result{}, transient, fmt.Errorf("API error: status %d: %.200s", resp.StatusCode, respBody)
	}
	if len(parsed.Choices) == 0 {
		return Result{}, false, errors.New("API response contained no choices")
	}
	msg := parsed.Choices[0].Message
	thinking, reply := SplitThinking(msg.Content)
	if msg.Reasoning == "" {
		msg.Reasoning = thinking
	}
	if msg.Reasoning != "" {
		emit(msg.Reasoning, true)
	}
	if reply != "" {
		emit(reply, false)
	}
	r = Result{Reply: reply, Reasoning: msg.Reasoning, Calls: msg.ToolCalls}
	if parsed.Usage.valid() {
		r.Usage = *parsed.Usage
	}
	if parsed.Choices[0].FinishReason == "length" {
		return r, false, c.truncated()
	}
	return r, false, nil
}

// SplitThinking separates reasoning a server left inline in the reply as
// <think>...</think>, rather than sending it separately, from the reply.
// Some chat templates open the <think> tag in the prompt, so only the
// closing tag may appear.
func SplitThinking(reply string) (thinking, answer string) {
	if i := strings.LastIndex(reply, "</think>"); i >= 0 {
		thinking = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(reply[:i]), "<think>"))
		return thinking, strings.TrimLeft(reply[i+len("</think>"):], " \t\r\n")
	}
	if t, ok := strings.CutPrefix(strings.TrimSpace(reply), "<think>"); ok {
		return strings.TrimSpace(t), "" // cut off while still thinking
	}
	return "", reply
}

// truncated builds the error for a reply cut off at the length limit.
func (c *Client) truncated() error {
	return fmt.Errorf("%w of %d tokens (raise max_tokens in agent.json, or with /config, if the model needs more room)", ErrTruncated, c.MaxTokens)
}

// streamChunk is one server-sent event of a streaming chat completion.
type streamChunk struct {
	Choices []struct {
		Delta struct {
			Content          string `json:"content"`
			ReasoningContent string `json:"reasoning_content"` // vLLM, llama.cpp, SGLang
			Reasoning        string `json:"reasoning"`         // the field name other servers use
			// Tool calls arrive in pieces: the first names the call, and
			// later ones with the same index add to its arguments.
			ToolCalls []struct {
				Index    int    `json:"index"`
				ID       string `json:"id"`
				Function struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `json:"function"`
			} `json:"tool_calls"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Usage *Usage `json:"usage"` // in the last chunk, if asked for
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// inlineSplitter routes reasoning a server leaves inline in the content, as
// <think>...</think>, to the reasoning side of the stream as it arrives.
type inlineSplitter struct {
	emit    func(string, bool)
	pending strings.Builder // content not yet known to be reply or reasoning
	state   splitState
}

type splitState int

const (
	atStart splitState = iota
	inThink
	inReply
)

func (s *inlineSplitter) write(text string) {
	if s.state == inReply {
		s.emit(text, false)
		return
	}
	s.pending.WriteString(text)
	buf := s.pending.String()
	if s.state == atStart {
		trimmed := strings.TrimLeft(buf, " \t\r\n")
		switch {
		case strings.HasPrefix(trimmed, "<think>"):
			s.state = inThink
			buf = strings.TrimPrefix(trimmed, "<think>")
		case len(trimmed) < len("<think>") && strings.HasPrefix("<think>", trimmed):
			return // could still become <think>
		case strings.Contains(buf, "</think>"):
			s.state = inThink // the template opened <think> in the prompt
		default:
			if trimmed == "" {
				return
			}
			s.state = inReply
			s.pending.Reset()
			s.emit(buf, false)
			return
		}
	}
	// Inside reasoning: emit all but a possible partial "</think>".
	if thought, rest, ok := strings.Cut(buf, "</think>"); ok {
		if thought != "" {
			s.emit(thought, true)
		}
		s.state = inReply
		s.pending.Reset()
		if rest = strings.TrimLeft(rest, " \t\r\n"); rest != "" {
			s.emit(rest, false)
		}
		return
	}
	keep := 0
	for k := len("</think>") - 1; k > 0; k-- {
		if strings.HasSuffix(buf, "</think>"[:k]) {
			keep = k
			break
		}
	}
	if len(buf) > keep {
		s.emit(buf[:len(buf)-keep], true)
	}
	s.pending.Reset()
	s.pending.WriteString(buf[len(buf)-keep:])
}

// readStream consumes a server-sent-event response, passing each token to
// emit as it arrives, and returns the assembled reply.
func (c *Client) readStream(resp *http.Response, start time.Time, emit func(string, bool)) (r Result, retryable bool, err error) {
	var logOut io.Writer = io.Discard
	if c.Logger != nil {
		head, err := httputil.DumpResponse(resp, false)
		if err != nil {
			head = []byte("dump failed: " + err.Error() + "\n")
		}
		c.Logger.Printf("response streaming after %v:\n%s", time.Since(start).Round(time.Millisecond), head)
		logOut = c.Logger.Writer()
	}
	var content, thought strings.Builder
	var calls []ToolCall
	var usage Usage
	separate := false // the server sends reasoning in its own field
	out := func(text string, reasoning bool) {
		fmt.Fprint(logOut, text)
		emit(text, reasoning)
	}
	split := &inlineSplitter{emit: out}
	generated := 0 // bytes of reply, reasoning and calls, for the local length limits
	limit := maxReplyBytes
	if c.MaxTokens > 0 {
		limit = min(limit, c.MaxTokens*bytesPerToken)
	}
	finished := false // [DONE] or a finish_reason came: the reply is whole
	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 0, 64<<10), 4<<20)
	for sc.Scan() {
		data, ok := strings.CutPrefix(sc.Text(), "data:")
		if !ok {
			continue // comment or blank separator line
		}
		data = strings.TrimSpace(data)
		if data == "[DONE]" {
			finished = true
			break
		}
		var chunk streamChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue // not a completion chunk (e.g. a keep-alive event)
		}
		if chunk.Error != nil {
			fmt.Fprintln(logOut)
			return Result{}, true, fmt.Errorf("API error mid-stream: %s", chunk.Error.Message)
		}
		if chunk.Usage.valid() {
			usage = *chunk.Usage
		}
		if len(chunk.Choices) == 0 {
			continue
		}
		delta := chunk.Choices[0].Delta
		if delta.ReasoningContent == "" {
			delta.ReasoningContent = delta.Reasoning
		}
		if delta.ReasoningContent != "" {
			separate = true
			out(delta.ReasoningContent, true)
			thought.WriteString(delta.ReasoningContent)
		}
		generated += len(delta.ReasoningContent) + len(delta.Content)
		for _, tc := range delta.ToolCalls {
			generated += len(tc.ID) + len(tc.Function.Name) + len(tc.Function.Arguments)
			if tc.Index < 0 || tc.Index > 128 {
				continue
			}
			for len(calls) <= tc.Index {
				calls = append(calls, ToolCall{Type: "function"})
			}
			call := &calls[tc.Index]
			if tc.ID != "" {
				call.ID = tc.ID
			}
			call.Function.Name += tc.Function.Name
			call.Function.Arguments += tc.Function.Arguments
		}
		if delta.Content != "" {
			if separate {
				out(delta.Content, false)
			} else {
				split.write(delta.Content)
			}
			content.WriteString(delta.Content)
		}
		if chunk.Choices[0].FinishReason != "" {
			finished = true
		}
		// Returning closes the response body, which stops generation on the
		// server too.
		if chunk.Choices[0].FinishReason == "length" || generated > limit {
			fmt.Fprintln(logOut)
			thinking, reply := SplitThinking(content.String())
			return Result{Reply: reply, Reasoning: thought.String() + thinking}, false, c.truncated()
		}
	}
	fmt.Fprintln(logOut)
	if err := sc.Err(); err != nil {
		return Result{}, true, fmt.Errorf("reading stream: %w", err)
	}
	if !finished {
		// Cut off without saying so: its tool calls may be cut short too.
		return Result{}, true, errors.New("the stream ended before the reply did")
	}
	if c.Logger != nil {
		c.Logger.Printf("response complete in %v", time.Since(start).Round(time.Millisecond))
	}
	thinking, reply := SplitThinking(content.String())
	return Result{Reply: reply, Reasoning: thought.String() + thinking, Calls: calls, Usage: usage}, false, nil
}

// redact hides the API key in logged output.
func redact(dump []byte, apiKey string) []byte {
	if apiKey == "" {
		return dump
	}
	return bytes.ReplaceAll(dump, []byte(apiKey), []byte("[redacted]"))
}

// TLSConfig builds the TLS settings for a server using a custom CA or
// self-signed certificate.
func TLSConfig(caCertPath string) (*tls.Config, error) {
	if caCertPath == "" {
		return &tls.Config{}, nil
	}
	pem, err := os.ReadFile(caCertPath)
	if err != nil {
		return nil, fmt.Errorf("reading ca-cert: %w", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(pem) {
		return nil, fmt.Errorf("no certificates found in %s", caCertPath)
	}
	return &tls.Config{RootCAs: pool}, nil
}