package models

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/go-resty/resty/v2"
)

// OpenAI-compatible Request/Response Structs

type ChatCompletionRequest struct {
	Model            string          `json:"model"`
	Messages         []ChatMessage   `json:"messages"`
	Temperature      *float64        `json:"temperature,omitempty"`
	TopP             *float64        `json:"top_p,omitempty"`
	TopK             *uint32         `json:"top_k,omitempty"`
	N                *uint32         `json:"n,omitempty"`
	Stop             json.RawMessage `json:"stop,omitempty"` // Can be string or array of strings
	MaxTokens        *uint32         `json:"max_tokens,omitempty"`
	PresencePenalty  *float64        `json:"presence_penalty,omitempty"`
	FrequencyPenalty *float64        `json:"frequency_penalty,omitempty"`
	Stream           *bool           `json:"stream,omitempty"`
	Tools            []Tool          `json:"tools,omitempty"`
	ToolChoice       json.RawMessage `json:"tool_choice,omitempty"`
}

type ChatMessage struct {
	Role             string      `json:"role"` // "system", "user", "assistant", "tool"
	Content          interface{} `json:"content"`
	ReasoningContent *string     `json:"reasoning_content,omitempty"`
	ToolCalls        []ToolCall  `json:"tool_calls,omitempty"`
	Name             *string     `json:"name,omitempty"`
	ToolCallID       *string     `json:"tool_call_id,omitempty"`
}

type ContentPart struct {
	Type     string    `json:"type"`
	Text     string    `json:"text,omitempty"`
	ImageUrl *ImageUrl `json:"image_url,omitempty"`
}

type ImageUrl struct {
	URL    string  `json:"url"`
	Detail *string `json:"detail,omitempty"`
}

// Tool-Calling Types

type Tool struct {
	Type     string      `json:"type"`
	Function FunctionDef `json:"function"`
}

type FunctionDef struct {
	Name        string          `json:"name"`
	Description *string         `json:"description,omitempty"`
	Parameters  json.RawMessage `json:"parameters,omitempty"`
}

type ToolCall struct {
	ID       string       `json:"id"`
	Type     string       `json:"type"`
	Function FunctionCall `json:"function"`
}

type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

type ToolCallDelta struct {
	Index    uint32             `json:"index"`
	ID       *string            `json:"id,omitempty"`
	Type     *string            `json:"type,omitempty"`
	Function *FunctionCallDelta `json:"function,omitempty"`
}

type FunctionCallDelta struct {
	Name      *string `json:"name,omitempty"`
	Arguments *string `json:"arguments,omitempty"`
}

// OpenAI-compatible Response Structs

type ChatCompletionResponse struct {
	ID      string       `json:"id"`
	Object  string       `json:"object"`
	Created int64        `json:"created"`
	Model   string       `json:"model"`
	Choices []ChatChoice `json:"choices"`
	Usage   *Usage       `json:"usage,omitempty"`
}

type ChatChoice struct {
	Index        uint32      `json:"index"`
	Message      ChatMessage `json:"message"`
	FinishReason *string     `json:"finish_reason,omitempty"`
}

type Usage struct {
	PromptTokens     uint32  `json:"prompt_tokens"`
	CompletionTokens uint32  `json:"completion_tokens"`
	TotalTokens      uint32  `json:"total_tokens"`
	ReasoningTokens  *uint32 `json:"reasoning_tokens,omitempty"`
	CacheReadTokens  *uint32 `json:"cache_read_tokens,omitempty"`
	CacheWriteTokens *uint32 `json:"cache_write_tokens,omitempty"`
}

// Streaming Response Structs

type ChatCompletionStreamResponse struct {
	ID      string             `json:"id"`
	Object  string             `json:"object"`
	Created int64              `json:"created"`
	Model   string             `json:"model"`
	Choices []ChatStreamChoice `json:"choices"`
	Usage   *Usage             `json:"usage,omitempty"`
	Error   *string            `json:"error,omitempty"`
}

type ChatStreamChoice struct {
	Index        uint32          `json:"index"`
	Delta        ChatStreamDelta `json:"delta"`
	FinishReason *string         `json:"finish_reason,omitempty"`
}

type ChatStreamDelta struct {
	Role             *string         `json:"role,omitempty"`
	Content          *string         `json:"content,omitempty"`
	ReasoningContent *string         `json:"reasoning_content,omitempty"`
	ToolCalls        []ToolCallDelta `json:"tool_calls,omitempty"`
}

type StreamUsage struct {
	PromptTokens     uint32 `json:"prompt_tokens"`
	CompletionTokens uint32 `json:"completion_tokens"`
	TotalTokens      uint32 `json:"total_tokens"`
	ReasoningTokens  uint32 `json:"reasoning_tokens"`
	CacheReadTokens  uint32 `json:"cache_read_tokens"`
	CacheWriteTokens uint32 `json:"cache_write_tokens"`
}

// Unified Request Format (for internal use)

type UnifiedRequest struct {
	ClientID         string
	Model            string
	Messages         []UnifiedMessage
	Temperature      *float64
	TopP             *float64
	TopK             *uint32
	N                *uint32
	Stop             []string
	MaxTokens        *uint32
	PresencePenalty  *float64
	FrequencyPenalty *float64
	Stream           bool
	HasImages        bool
	Tools            []Tool
	ToolChoice       json.RawMessage
}

type UnifiedMessage struct {
	Role             string
	Content          []UnifiedContentPart
	ReasoningContent *string
	ToolCalls        []ToolCall
	Name             *string
	ToolCallID       *string
}

type UnifiedContentPart struct {
	Type  string
	Text  string
	Image *ImageInput
}

type ImageInput struct {
	Base64   string `json:"base64,omitempty"`
	URL      string `json:"url,omitempty"`
	MimeType string `json:"mime_type,omitempty"`
}

func (i *ImageInput) ToBase64() (string, string, error) {
	if i.Base64 != "" {
		return i.Base64, i.MimeType, nil
	}

	if i.URL != "" {
		client := resty.New()
		resp, err := client.R().Get(i.URL)
		if err != nil {
			return "", "", fmt.Errorf("failed to fetch image: %w", err)
		}

		if !resp.IsSuccess() {
			return "", "", fmt.Errorf("failed to fetch image: HTTP %d", resp.StatusCode())
		}

		mimeType := resp.Header().Get("Content-Type")
		if mimeType == "" {
			mimeType = "image/jpeg"
		}

		encoded := base64.StdEncoding.EncodeToString(resp.Body())
		return encoded, mimeType, nil
	}

	return "", "", fmt.Errorf("empty image input")
}

// Image Generation (DALL-E, Imagen)

type ImageGenerationRequest struct {
	Model          string  `json:"model"`
	Prompt         string  `json:"prompt"`
	N              *uint32 `json:"n,omitempty"`
	Quality        *string `json:"quality,omitempty"`
	ResponseFormat *string `json:"response_format,omitempty"`
	Size           *string `json:"size,omitempty"`
	Style          *string `json:"style,omitempty"`
	User           *string `json:"user,omitempty"`
}

type ImageGenerationResponse struct {
	Created int64       `json:"created"`
	Data    []ImageData `json:"data"`
}

type ImageData struct {
	URL           string `json:"url,omitempty"`
	B64JSON       string `json:"b64_json,omitempty"`
	RevisedPrompt string `json:"revised_prompt,omitempty"`
}

// AuthInfo for context
type AuthInfo struct {
	Token    string
	ClientID string
}
