package types

import (
	"encoding/json"
	"strings"
)

func StringPtr(s string) *string { return &s }

// Request
type ChatMessageContent string

type ChatMessageContentPart struct {
	Type string `json:"type"`
	Text string `json:"text,omitempty"`
}

func (c *ChatMessageContent) UnmarshalJSON(data []byte) error {
	var text string
	if err := json.Unmarshal(data, &text); err == nil {
		*c = ChatMessageContent(text)
		return nil
	}

	var parts []ChatMessageContentPart
	if err := json.Unmarshal(data, &parts); err != nil {
		return err
	}

	var content strings.Builder
	for _, part := range parts {
		if part.Type == "text" {
			content.WriteString(part.Text)
		}
	}
	*c = ChatMessageContent(content.String())
	return nil
}

type ChatMessage struct {
	Role    string             `json:"role"`
	Content ChatMessageContent `json:"content"`
}

type ChatCompletionRequest struct {
	Model       string        `json:"model"`
	Messages    []ChatMessage `json:"messages"`
	Stream      bool          `json:"stream"`
	Temperature *float64      `json:"temperature,omitempty"`
}

// Response (non-streaming)
type ChatCompletionResponse struct {
	ID      string   `json:"id"`
	Object  string   `json:"object"`
	Created int64    `json:"created"`
	Model   string   `json:"model"`
	Choices []Choice `json:"choices"`
	Usage   Usage    `json:"usage"`
}

type Choice struct {
	Index        int         `json:"index"`
	Message      ChatMessage `json:"message"`
	FinishReason *string     `json:"finish_reason"`
}

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// Streaming chunk
type ChatCompletionChunk struct {
	ID      string        `json:"id"`
	Object  string        `json:"object"`
	Created int64         `json:"created"`
	Model   string        `json:"model,omitempty"`
	Choices []ChunkChoice `json:"choices"`
	Usage   *Usage        `json:"usage,omitempty"`
}

type ChunkChoice struct {
	Index        int     `json:"index"`
	Delta        Delta   `json:"delta"`
	FinishReason *string `json:"finish_reason"`
}

type Delta struct {
	Role    *string `json:"role,omitempty"`
	Content *string `json:"content,omitempty"`
}

// Models list
type ModelList struct {
	Object string      `json:"object"`
	Data   []ModelInfo `json:"data"`
}

type ModelInfo struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

// Error response
type ErrorResponse struct {
	Error ErrorBody `json:"error"`
}

type ErrorBody struct {
	Message string `json:"message"`
	Type    string `json:"type"`
	Code    string `json:"code,omitempty"`
}

func NewErrorResponse(message, errType, code string) ErrorResponse {
	return ErrorResponse{
		Error: ErrorBody{
			Message: message,
			Type:    errType,
			Code:    code,
		},
	}
}

func NewChatCompletionChunk(id string, created int64, model string, delta Delta) ChatCompletionChunk {
	return ChatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []ChunkChoice{
			{
				Index: 0,
				Delta: delta,
			},
		},
	}
}
