package llm

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

// CompletionsParams contains the inputs to a chat completion.
type CompletionsParams struct {
	// Messages is the conversation presented to the model, in order.
	Messages []Message
	// Model identifies the model according to the configured endpoint.
	Model string
	// ThinkingBudgetTokens is sent as the endpoint's reasoning-token budget.
	// Nil omits both the budget and enable_thinking; zero sets enable_thinking
	// to false; any other value sets it to true.
	ThinkingBudgetTokens *int
	// Tools describes the functions available for the model to call.
	Tools []APITool
}

// CompletionResult contains one generated Message and its response metadata.
type CompletionResult struct {
	// FinishReason is the endpoint's reason for ending generation.
	FinishReason string
	// Message is the generated message.
	Message Message
	// Timings contains llama.cpp performance data, when reported.
	Timings *Timings
	// Usage contains token counts, when reported.
	Usage *Usage
}

// Completions requests a non-streaming chat completion and returns the first
// choice. It returns an error if the response contains no choices.
func (c *Client) Completions(
	ctx context.Context, p CompletionsParams,
) (*CompletionResult, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	cr, err := c.completions(ctx, p)
	if err != nil {
		return nil, err
	}
	choice := cr.Choices[0]
	return &CompletionResult{
		Message:      choice.Message,
		FinishReason: choice.FinishReason,
		Usage:        cr.Usage,
		Timings:      cr.Timings,
	}, nil
}

// completions sends a single non-streaming request and returns the
// full API response.
func (c *Client) completions(
	ctx context.Context, p CompletionsParams,
) (*ChatResponse, error) {
	kwargs := map[string]any{}
	if p.ThinkingBudgetTokens != nil {
		kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
	}
	body, err := json.Marshal(ChatRequest{
		Model:                p.Model,
		Messages:             p.Messages,
		Tools:                p.Tools,
		ThinkingBudgetTokens: p.ThinkingBudgetTokens,
		ChatTemplateKWARGS:   kwargs,
		ParallelToolCalls:    new(true),
	})
	if err != nil {
		return nil, fmt.Errorf(
			"completions: marshal request body: %w", err,
		)
	}
	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		c.baseURL+chatCompletionsPath,
		bytes.NewReader(body),
	)
	if err != nil {
		return nil, fmt.Errorf("completions: create request: %w", err)
	}
	if c.key != "" {
		req.Header.Set("Authorization", "Bearer "+c.key)
	}
	req.Header.Set("Content-Type", "application/json")

	res, err := c.httpc.Do(req)
	if err != nil {
		return nil, fmt.Errorf("completions: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		body, err := io.ReadAll(res.Body)
		if err != nil {
			c.log.ErrorContext(
				ctx,
				"completions: failed to read error response body",
				slog.Any("error", err),
			)
			return nil, fmt.Errorf(
				"completions: status code %d", res.StatusCode,
			)
		}
		var apiError struct {
			Error *APIError `json:"error"`
		}
		if err := json.Unmarshal(
			body, &apiError,
		); err == nil && apiError.Error != nil {
			c.log.ErrorContext(
				ctx,
				"completions: request error",
				slog.Int("status", res.StatusCode),
				slog.String("code", apiError.Error.Code),
				slog.String(
					"message", apiError.Error.Message,
				),
				slog.String("type", apiError.Error.Type),
			)
			return nil, fmt.Errorf("completions: %w", apiError.Error)
		}
		c.log.ErrorContext(
			ctx,
			"completions: API error",
			slog.Int("status", res.StatusCode),
			slog.String("body", string(body)),
		)
		return nil, fmt.Errorf(
			"completions: %d: %s",
			res.StatusCode, string(body),
		)
	}

	var cr = &ChatResponse{}
	if err := json.NewDecoder(res.Body).Decode(cr); err != nil {
		return nil, fmt.Errorf("completions: decode response: %w", err)
	}
	if len(cr.Choices) == 0 {
		return nil, fmt.Errorf(
			"completions: no response choices returned",
		)
	}
	return cr, nil
}

// StreamAccumulator reconstructs a Message from the ordered deltas for one
// streaming choice. The zero value is ready for use.
type StreamAccumulator struct {
	// Content contains the accumulated message text.
	Content strings.Builder
	// Reasoning combines fragments received under either supported reasoning
	// field name.
	Reasoning strings.Builder
	// Refusal contains the accumulated refusal text.
	Refusal strings.Builder
	// Role is the latest non-empty role received.
	Role Role
	// ToolCalls contains calls assembled by tool-call index.
	ToolCalls []ToolCall
}

// Add incorporates the next delta. Content, refusal, reasoning, and tool-call
// fragments are appended in arrival order. Tool-call fragments are grouped by
// Index; a missing index denotes the first call. A non-empty Role replaces the
// previously accumulated role.
func (acc *StreamAccumulator) Add(delta Delta) {
	if delta.Role != "" {
		acc.Role = delta.Role
	}
	acc.Content.WriteString(delta.Content)
	acc.Refusal.WriteString(delta.Refusal)
	acc.Reasoning.WriteString(delta.ReasoningContent)
	acc.Reasoning.WriteString(delta.Reasoning)

	for _, tc := range delta.ToolCalls {
		idx := 0
		if tc.Index != nil {
			idx = *tc.Index
		}
		for len(acc.ToolCalls) <= idx {
			acc.ToolCalls = append(
				acc.ToolCalls,
				ToolCall{Type: ToolTypeFunction},
			)
		}
		if tc.ID != "" {
			acc.ToolCalls[idx].ID = tc.ID
		}
		if tc.Type != "" {
			acc.ToolCalls[idx].Type = tc.Type
		}
		acc.ToolCalls[idx].Function.Name += tc.Function.Name
		acc.ToolCalls[idx].Function.Arguments += tc.Function.Arguments
	}
}

// Message returns the accumulated state without resetting the accumulator. If
// the stream contains a refusal, the returned content consists of that refusal
// rather than the accumulated text.
func (acc *StreamAccumulator) Message() Message {
	var contentParts []ContentPart
	if refusal := acc.Refusal.String(); refusal != "" {
		contentParts = []ContentPart{{
			Type:    ContentTypeRefusal,
			Refusal: refusal,
		}}
	} else {
		contentParts = []ContentPart{{
			Type: ContentTypeText,
			Text: acc.Content.String(),
		}}
	}

	return Message{
		ContentParts:     contentParts,
		ReasoningContent: acc.Reasoning.String(),
		Refusal:          acc.Refusal.String(),
		Role:             acc.Role,
		ToolCalls:        acc.ToolCalls,
	}
}

// CompletionStream is a caller-driven chat-completion stream. A stream owns
// its response body and timeout context until it terminates or is closed.
type CompletionStream struct {
	ctx    context.Context
	reader *bufio.Reader
	body   io.ReadCloser
	cancel context.CancelFunc

	acc           StreamAccumulator
	current       StreamChunk
	finishReason  string
	usage         *Usage
	timings       *Timings
	err           error
	cleanupErr    error
	terminated    bool
	resourcesDone bool
}

// CompletionsStream opens a streaming chat completion. Request setup,
// transport, and non-success response errors are returned before a stream is
// created.
func (c *Client) CompletionsStream(
	ctx context.Context,
	p CompletionsParams,
) (*CompletionStream, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)

	kwargs := map[string]any{}
	if p.ThinkingBudgetTokens != nil {
		kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
	}
	body, err := json.Marshal(ChatRequest{
		ChatTemplateKWARGS:   kwargs,
		Messages:             p.Messages,
		Model:                p.Model,
		ParallelToolCalls:    new(true),
		Stream:               true,
		StreamOptions:        &StreamOptions{IncludeUsage: true},
		ThinkingBudgetTokens: p.ThinkingBudgetTokens,
		Tools:                p.Tools,
	})
	if err != nil {
		cancel()
		return nil, fmt.Errorf(
			"completions stream: marshal request body: %w", err,
		)
	}

	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		c.baseURL+chatCompletionsPath,
		bytes.NewReader(body),
	)
	if err != nil {
		cancel()
		return nil, fmt.Errorf(
			"completions stream: create request: %w", err,
		)
	}
	if c.key != "" {
		req.Header.Set("Authorization", "Bearer "+c.key)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")

	res, err := c.httpc.Do(req)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("completions stream: %w", err)
	}

	if res.StatusCode != http.StatusOK {
		defer cancel()
		defer res.Body.Close()
		body, err := io.ReadAll(res.Body)
		if err != nil {
			c.log.ErrorContext(
				ctx,
				"completions stream: failed to read error response body",
				slog.Any("error", err),
			)
			return nil, fmt.Errorf(
				"completions stream: status code %d",
				res.StatusCode,
			)
		}
		var apiError struct {
			Error *APIError `json:"error"`
		}
		if err := json.Unmarshal(
			body, &apiError,
		); err == nil && apiError.Error != nil {
			c.log.ErrorContext(
				ctx,
				"completions stream: request error",
				slog.Int("status", res.StatusCode),
				slog.String("code", apiError.Error.Code),
				slog.String("message", apiError.Error.Message),
				slog.String("type", apiError.Error.Type),
			)
			return nil, fmt.Errorf(
				"completions stream: %w", apiError.Error,
			)
		}
		c.log.ErrorContext(
			ctx,
			"completions stream: API error",
			slog.Int("status", res.StatusCode),
			slog.String("body", string(body)),
		)
		return nil, fmt.Errorf(
			"completions stream: %d: %s",
			res.StatusCode,
			string(body),
		)
	}

	return &CompletionStream{
		ctx:    ctx,
		reader: bufio.NewReader(res.Body),
		body:   res.Body,
		cancel: cancel,
	}, nil
}

// Next advances to the next JSON data event. It returns false after [DONE], an
// error, or Close. On an error, Err reports the cause and Result remains
// available with all chunks accumulated before the error.
func (s *CompletionStream) Next() bool {
	if s.terminated {
		return false
	}

	var data []string
	for {
		line, err := s.reader.ReadString('\n')
		if err != nil && err != io.EOF {
			s.fail(fmt.Errorf("completions stream: read response: %w", err))
			return false
		}
		if len(line) > 0 {
			line = strings.TrimSuffix(line, "\n")
			line = strings.TrimSuffix(line, "\r")
			if line == "" {
				if len(data) > 0 {
					break
				}
			} else if line[0] != ':' {
				field, value, found := strings.Cut(line, ":")
				if !found {
					field = line
					value = ""
				}
				value = strings.TrimPrefix(value, " ")
				if field == "data" {
					data = append(data, value)
				}
			}
		}
		if err == io.EOF {
			s.fail(io.ErrUnexpectedEOF)
			return false
		}
	}

	payload := strings.Join(data, "\n")
	if payload == sseDone {
		s.terminated = true
		s.release()
		return false
	}

	var chunk StreamChunk
	if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
		s.fail(fmt.Errorf("completions stream: decode chunk: %w", err))
		return false
	}
	for _, choice := range chunk.Choices {
		if choice.Index != 0 {
			s.fail(fmt.Errorf(
				"completions stream: unexpected choice index %d",
				choice.Index,
			))
			return false
		}
		if choice.FinishReason != "" {
			s.finishReason = choice.FinishReason
		}
		s.acc.Add(choice.Delta)
	}
	if chunk.Usage != nil {
		s.usage = chunk.Usage
	}
	if chunk.Timings != nil {
		s.timings = chunk.Timings
	}
	s.current = chunk
	return true
}

// Current returns the most recently decoded chunk. The value is replaced by
// the next successful call to Next.
func (s *CompletionStream) Current() StreamChunk {
	return s.current
}

// Result returns a snapshot of the state accumulated so far. Err reports
// whether the snapshot is a complete response or a partial result.
func (s *CompletionStream) Result() CompletionResult {
	message := s.acc.Message()
	message.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
	return CompletionResult{
		Message:      message,
		FinishReason: s.finishReason,
		Usage:        s.usage,
		Timings:      s.timings,
	}
}

// Err returns the stream failure, if any. Closing a stream does not create or
// replace a stream failure.
func (s *CompletionStream) Err() error {
	return s.err
}

// Close releases the response body and timeout context. It is safe to call
// repeatedly and returns the first response-body cleanup error.
func (s *CompletionStream) Close() error {
	s.terminated = true
	s.release()
	return s.cleanupErr
}

func (s *CompletionStream) fail(err error) {
	if ctxErr := s.ctx.Err(); ctxErr != nil {
		err = ctxErr
	}
	s.err = err
	s.terminated = true
	s.release()
}

func (s *CompletionStream) release() {
	if s.resourcesDone {
		return
	}
	s.resourcesDone = true
	if err := s.body.Close(); err != nil {
		s.cleanupErr = err
	}
	s.cancel()
}