// Package sglang provides a Go SDK for SGLang gRPC API.
//
// SGLang is a fast language model serving framework. This package provides a Go client
// library for interacting with SGLang's gRPC API, following the style of OpenAI's Go SDK.
//
// Basic usage:
//
//	client, err := sglang.NewClient(sglang.ClientConfig{
//		Endpoint:      "grpc://localhost:20000",
//		TokenizerPath: "/path/to/tokenizer",
//	})
//	if err != nil {
//		log.Fatal(err)
//	}
//	defer client.Close()
//
//	resp, err := client.CreateChatCompletion(ctx, sglang.ChatCompletionRequest{
//		Model: "default",
//		Messages: []sglang.ChatMessage{
//			{Role: "user", Content: "Hello"},
//		},
//	})
//
// For streaming responses, use CreateChatCompletionStream instead.
package sglang

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	grpcclient "github.com/sglang/sglang-go-grpc-sdk/internal/grpc"
)

// Client is the main client for interacting with SGLang gRPC API.
// It manages the connection to the SGLang server and handles both streaming
// and non-streaming chat completions.
//
// Thread-safe: All public methods are safe for concurrent use.
type Client struct {
	endpoint      string
	tokenizerPath string
	grpcClient    *grpcclient.GrpcClient // gRPC-based client
	mu            sync.RWMutex
}

// ClientConfig holds configuration for creating a new client.
type ClientConfig struct {
	// Endpoint is the gRPC endpoint URL (e.g., "grpc://localhost:20000").
	// Required field. Must include the scheme (grpc://) and port number.
	Endpoint string

	// TokenizerPath is the path to the tokenizer directory containing
	// tokenizer configuration files (e.g., tokenizer.json, vocab.json).
	// Required field.
	TokenizerPath string

	// ChannelBufferSizes configures buffer sizes for internal channels.
	// If nil, default values will be used (optimized for high concurrency).
	ChannelBufferSizes *ChannelBufferSizes

	// Timeouts configures timeout values for various operations.
	// If nil, default values will be used.
	Timeouts *Timeouts
}

// ChannelBufferSizes configures buffer sizes for internal channels.
// These affect concurrency and memory usage. Larger buffers allow more
// concurrent operations but use more memory.
type ChannelBufferSizes = grpcclient.ChannelBufferSizes

// Timeouts configures timeout values for various operations.
type Timeouts = grpcclient.Timeouts

// defaultChannelBufferSizes returns default channel buffer sizes optimized for high concurrency (10k+).
// These values are designed to handle thousands of concurrent requests without blocking.
func defaultChannelBufferSizes() ChannelBufferSizes {
	return ChannelBufferSizes{
		ResultJSONChan: 10000, // Increased for high concurrency: each request may produce 200-500 chunks
		ErrChan:        100,   // Errors are rare, 100 is sufficient
		RecvChan:       2000,  // Increased for high concurrency: more gRPC responses to buffer
	}
}

// defaultTimeouts returns default timeout values.
func defaultTimeouts() Timeouts {
	return Timeouts{
		KeepaliveTime:    300 * time.Second, // Increased to reduce ping frequency and avoid "too many pings" errors
		KeepaliveTimeout: 20 * time.Second,
		CloseTimeout:     5 * time.Second,
	}
}

// NewClient creates a new SGLang client with the given configuration.
//
// The client maintains a long-lived connection to the SGLang server and should
// be reused for multiple requests. Call Close() to release resources.
//
// Returns an error if:
// - Endpoint is empty
// - TokenizerPath is empty
// - Connection to the server fails
func NewClient(config ClientConfig) (*Client, error) {
	if config.Endpoint == "" {
		return nil, errors.New("endpoint is required")
	}
	if config.TokenizerPath == "" {
		return nil, errors.New("tokenizer path is required")
	}

	bufferSizes := defaultChannelBufferSizes()
	if config.ChannelBufferSizes != nil {
		if config.ChannelBufferSizes.ResultJSONChan > 0 {
			bufferSizes.ResultJSONChan = config.ChannelBufferSizes.ResultJSONChan
		}
		if config.ChannelBufferSizes.ErrChan > 0 {
			bufferSizes.ErrChan = config.ChannelBufferSizes.ErrChan
		}
		if config.ChannelBufferSizes.RecvChan > 0 {
			bufferSizes.RecvChan = config.ChannelBufferSizes.RecvChan
		}
	}

	timeouts := defaultTimeouts()
	if config.Timeouts != nil {
		if config.Timeouts.KeepaliveTime > 0 {
			timeouts.KeepaliveTime = config.Timeouts.KeepaliveTime
		}
		if config.Timeouts.KeepaliveTimeout > 0 {
			timeouts.KeepaliveTimeout = config.Timeouts.KeepaliveTimeout
		}
		if config.Timeouts.CloseTimeout > 0 {
			timeouts.CloseTimeout = config.Timeouts.CloseTimeout
		}
	}

	grpcClient, err := grpcclient.NewGrpcClient(config.Endpoint, config.TokenizerPath, bufferSizes, timeouts)
	if err != nil {
		return nil, fmt.Errorf("failed to create gRPC client: %w", err)
	}

	return &Client{
		endpoint:      config.Endpoint,
		tokenizerPath: config.TokenizerPath,
		grpcClient:    grpcClient,
	}, nil
}

// Close closes the client and releases all resources.
//
// After Close() is called, the client cannot be used for further requests.
// Calling Close() multiple times is safe and idempotent.
func (c *Client) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.grpcClient != nil {
		if err := c.grpcClient.Close(); err != nil {
			return err
		}
		c.grpcClient = nil
	}
	return nil
}

// ChatCompletionRequest represents a request for chat completion.
// It follows the OpenAI API style for familiar usage.
type ChatCompletionRequest struct {
	// Model specifies the model to use for completion (e.g., "default")
	Model string `json:"model"`
	// Messages is the list of messages in the conversation
	Messages            []ChatMessage   `json:"messages"`
	Temperature         *float32        `json:"temperature,omitempty"`
	TopP                *float32        `json:"top_p,omitempty"`
	TopK                *int            `json:"top_k,omitempty"`
	MaxCompletionTokens *int            `json:"max_completion_tokens,omitempty"`
	Stream              bool            `json:"stream"`
	Tools               []Tool          `json:"tools,omitempty"`
	ToolChoice          interface{}     `json:"tool_choice,omitempty"`
	Stop                interface{}     `json:"stop,omitempty"`
	StopTokenIDs        []int           `json:"stop_token_ids,omitempty"`
	SkipSpecialTokens   bool            `json:"skip_special_tokens,omitempty"`
	FrequencyPenalty    *float32        `json:"frequency_penalty,omitempty"`
	PresencePenalty     *float32        `json:"presence_penalty,omitempty"`
	ResponseFormat      *ResponseFormat `json:"response_format,omitempty"`
	Seed                *int            `json:"seed,omitempty"`
	Logprobs            bool            `json:"logprobs,omitempty"`
	TopLogprobs         *int            `json:"top_logprobs,omitempty"`
	User                string          `json:"user,omitempty"`
}

// ChatMessage represents a single message in a chat conversation
type ChatMessage struct {
	Role    string      `json:"role"`
	Content interface{} `json:"content"`
	Name    string      `json:"name,omitempty"`
}

// Tool represents a tool/function that can be called
type Tool struct {
	Type     string   `json:"type"`
	Function Function `json:"function"`
}

// Function represents a function definition
type Function struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description,omitempty"`
	Parameters  map[string]interface{} `json:"parameters"`
}

// ResponseFormat represents the response format
type ResponseFormat struct {
	Type string `json:"type"`
}

// ChatCompletionResponse represents a non-streaming chat completion response
type ChatCompletionResponse struct {
	ID                string   `json:"id"`
	Object            string   `json:"object"`
	Created           int64    `json:"created"`
	Model             string   `json:"model"`
	SystemFingerprint string   `json:"system_fingerprint,omitempty"`
	Choices           []Choice `json:"choices"`
	Usage             Usage    `json:"usage"`
}

// Choice represents a choice in the completion response
type Choice struct {
	Index        int     `json:"index"`
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

// Message represents a message in the response
type Message struct {
	Role      string     `json:"role"`
	Content   string     `json:"content"`
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

// ToolCall represents a tool call in the response
type ToolCall struct {
	ID       string       `json:"id"`
	Type     string       `json:"type"`
	Function FunctionCall `json:"function"`
}

// FunctionCall represents a function call
type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// Usage represents token usage information
type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// ChatCompletionStreamResponse represents a streaming chat completion response
type ChatCompletionStreamResponse struct {
	ID                string         `json:"id"`
	Object            string         `json:"object"`
	Created           int64          `json:"created"`
	Model             string         `json:"model"`
	SystemFingerprint string         `json:"system_fingerprint,omitempty"`
	Choices           []StreamChoice `json:"choices"`
	Usage             *Usage         `json:"usage,omitempty"`
}

// StreamChoice represents a choice in a streaming response
type StreamChoice struct {
	Index        int          `json:"index"`
	Delta        MessageDelta `json:"delta"`
	FinishReason string       `json:"finish_reason,omitempty"`
}

// MessageDelta represents incremental message updates
type MessageDelta struct {
	Role      string     `json:"role,omitempty"`
	Content   string     `json:"content,omitempty"`
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

// CreateChatCompletion creates a non-streaming chat completion with context support.
//
// Context Support:
// The ctx parameter is fully supported for cancellation and timeouts:
// - If ctx is cancelled, the request will be interrupted on the next stream.RecvJSON() call
// - If ctx times out, the request will return context.DeadlineExceeded
//
// Example with timeout:
//
//	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
//	defer cancel()
//	resp, err := client.CreateChatCompletion(ctx, req)
//
// Note: Internally, this creates a stream and collects all chunks,
// so context monitoring happens at the chunk level.
func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) {
	// For non-streaming, we'll collect all chunks and return the final response
	req.Stream = true // We still use streaming internally, but collect all chunks

	if len(req.Tools) == 0 {
		req.Tools = nil
	}

	stream, err := c.CreateChatCompletionStream(ctx, req)
	if err != nil {
		return nil, err
	}
	defer stream.Close()

	var fullContent strings.Builder
	var fullToolCalls []ToolCall
	var finishReason string
	var usage Usage
	var responseID string
	var created int64
	var model string
	var systemFingerprint string

	for {
		chunkJSON, err := stream.RecvJSON()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		var chunk ChatCompletionStreamResponse
		if err := json.Unmarshal([]byte(chunkJSON), &chunk); err != nil {
			return nil, fmt.Errorf("failed to parse chunk: %w", err)
		}

		if chunk.ID != "" {
			responseID = chunk.ID
		}
		if chunk.Created > 0 {
			created = chunk.Created
		}
		if chunk.Model != "" {
			model = chunk.Model
		}
		if chunk.SystemFingerprint != "" {
			systemFingerprint = chunk.SystemFingerprint
		}

		for _, choice := range chunk.Choices {
			if choice.Delta.Content != "" {
				fullContent.WriteString(choice.Delta.Content)
			}
			if len(choice.Delta.ToolCalls) > 0 {
				fullToolCalls = append(fullToolCalls, choice.Delta.ToolCalls...)
			}
			if choice.FinishReason != "" {
				finishReason = choice.FinishReason
			}
		}

		if chunk.Usage != nil {
			usage = *chunk.Usage
		}
	}

	message := Message{
		Role:    "assistant",
		Content: fullContent.String(),
	}
	if len(fullToolCalls) > 0 {
		message.ToolCalls = fullToolCalls
	}

	if finishReason == "" {
		finishReason = "stop"
	}

	return &ChatCompletionResponse{
		ID:                responseID,
		Object:            "chat.completion",
		Created:           created,
		Model:             model,
		SystemFingerprint: systemFingerprint,
		Choices: []Choice{
			{
				Index:        0,
				Message:      message,
				FinishReason: finishReason,
			},
		},
		Usage: usage,
	}, nil
}

// ChatCompletionStream represents a streaming chat completion
type ChatCompletionStream struct {
	grpcStream *grpcclient.GrpcChatCompletionStream
	ctx        context.Context
	cancel     context.CancelFunc
}

func (s *ChatCompletionStream) RecvJSON() (string, error) {
	return s.grpcStream.RecvJSON()
}

// Close closes the stream and cancels any pending operations.
func (s *ChatCompletionStream) Close() error {
	if s.cancel != nil {
		s.cancel()
	}
	if s.grpcStream != nil {
		return s.grpcStream.Close()
	}
	return nil
}

// CreateChatCompletionStream creates a streaming chat completion with context cancellation support.
//
// Context Support:
// The ctx parameter is now fully supported for cancellation and timeouts:
// - If ctx is cancelled, stream.RecvJSON() will return context.Canceled on the next call
// - If ctx times out (WithTimeout), stream.RecvJSON() will return context.DeadlineExceeded
// - Calling stream.Close() also cancels the context
//
// Example with timeout:
//
//	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
//	defer cancel()
//	stream, err := client.CreateChatCompletionStream(ctx, req)
//	// Stream will auto-close if 30 seconds elapse
//
// Example with cancellation:
//
//	ctx, cancel := context.WithCancel(context.Background())
//	stream, err := client.CreateChatCompletionStream(ctx, req)
//	go func() {
//	    time.Sleep(5*time.Second)
//	    cancel()  // Cancel after 5 seconds
//	}()
func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionStream, error) {
	reqJSON, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	var reqMap map[string]interface{}
	if err := json.Unmarshal(reqJSON, &reqMap); err != nil {
		return nil, fmt.Errorf("failed to unmarshal request to map: %w", err)
	}

	if _, exists := reqMap["tools"]; !exists {
		reqMap["tools"] = []interface{}{}
	}

	reqJSON, err = json.Marshal(reqMap)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request map to JSON: %w", err)
	}

	if c.grpcClient == nil {
		return nil, errors.New("gRPC client is closed")
	}

	grpcStream, err := c.grpcClient.CreateChatCompletionStream(ctx, string(reqJSON))
	if err != nil {
		return nil, fmt.Errorf("failed to create gRPC stream: %w", err)
	}

	streamCtx, cancel := context.WithCancel(ctx)
	return &ChatCompletionStream{
		grpcStream: grpcStream,
		ctx:        streamCtx,
		cancel:     cancel,
	}, nil
}