package ollama

import (
	"time"
)

// GenerateRequest represents a request to generate text
type GenerateRequest struct {
	Model     string                 `json:"model" validate:"required"`
	Prompt    string                 `json:"prompt" validate:"required"`
	Stream    *bool                  `json:"stream,omitempty"`
	Format    string                 `json:"format,omitempty"`
	Options   map[string]interface{} `json:"options,omitempty"`
	System    string                 `json:"system,omitempty"`
	Template  string                 `json:"template,omitempty"`
	Context   []int                  `json:"context,omitempty"`
	Raw       bool                   `json:"raw,omitempty"`
	KeepAlive string                 `json:"keep_alive,omitempty"`
}

// GenerateResponse represents the response from a text generation request
type GenerateResponse struct {
	Model              string    `json:"model"`
	CreatedAt          time.Time `json:"created_at"`
	Response           string    `json:"response"`
	Done               bool      `json:"done"`
	Context            []int     `json:"context,omitempty"`
	TotalDuration      int64     `json:"total_duration,omitempty"`
	LoadDuration       int64     `json:"load_duration,omitempty"`
	PromptEvalCount    int       `json:"prompt_eval_count,omitempty"`
	PromptEvalDuration int64     `json:"prompt_eval_duration,omitempty"`
	EvalCount          int       `json:"eval_count,omitempty"`
	EvalDuration       int64     `json:"eval_duration,omitempty"`
}

// ChatRequest represents a request for a chat conversation
type ChatRequest struct {
	Model     string      `json:"model" validate:"required"`
	Messages  []Message   `json:"messages" validate:"required,min=1"`
	Stream    *bool       `json:"stream,omitempty"`
	Format    string      `json:"format,omitempty"`
	Options   ChatOptions `json:"options,omitempty"`
	KeepAlive string      `json:"keep_alive,omitempty"`
}

// Message represents a message in a chat conversation
type Message struct {
	Role    MessageRole `json:"role" validate:"required,oneof=system user assistant"`
	Content string      `json:"content" validate:"required"`
	Images  []string    `json:"images,omitempty"`
}

// ChatOptions represents options for chat requests
type ChatOptions struct {
	Temperature   *float64               `json:"temperature,omitempty"`
	TopP          *float64               `json:"top_p,omitempty"`
	TopK          *int                   `json:"top_k,omitempty"`
	RepeatPenalty *float64               `json:"repeat_penalty,omitempty"`
	Seed          *int                   `json:"seed,omitempty"`
	Stop          []string               `json:"stop,omitempty"`
	NumCtx        *int                   `json:"num_ctx,omitempty"`
	NumPredict    *int                   `json:"num_predict,omitempty"`
	NumKeep       *int                   `json:"num_keep,omitempty"`
	NumThread     *int                   `json:"num_thread,omitempty"`
	RepeatLastN   *int                   `json:"repeat_last_n,omitempty"`
	Options       map[string]interface{} `json:"options,omitempty"`
}

// ChatResponse represents the response from a chat request
type ChatResponse struct {
	Model              string    `json:"model"`
	CreatedAt          time.Time `json:"created_at"`
	Message            Message   `json:"message"`
	Done               bool      `json:"done"`
	TotalDuration      int64     `json:"total_duration,omitempty"`
	LoadDuration       int64     `json:"load_duration,omitempty"`
	PromptEvalCount    int       `json:"prompt_eval_count,omitempty"`
	PromptEvalDuration int64     `json:"prompt_eval_duration,omitempty"`
	EvalCount          int       `json:"eval_count,omitempty"`
	EvalDuration       int64     `json:"eval_duration,omitempty"`
}

// EmbedRequest represents a request to generate embeddings
type EmbedRequest struct {
	Model   string                 `json:"model" validate:"required"`
	Prompt  string                 `json:"prompt" validate:"required"`
	Options map[string]interface{} `json:"options,omitempty"`
}

// EmbedResponse represents the response from an embedding request
type EmbedResponse struct {
	Embedding []float64 `json:"embedding"`
}

// PullRequest represents a request to pull a model
type PullRequest struct {
	Name     string `json:"name" validate:"required"`
	Insecure bool   `json:"insecure,omitempty"`
	Stream   *bool  `json:"stream,omitempty"`
}

// PullResponse represents the response from a pull request
type PullResponse struct {
	Status    string `json:"status"`
	Digest    string `json:"digest,omitempty"`
	Total     int64  `json:"total,omitempty"`
	Completed int64  `json:"completed,omitempty"`
}

// ListModelsResponse represents the response from listing models
type ListModelsResponse struct {
	Models []ModelInfo `json:"models"`
}

// ModelInfo represents information about a model
type ModelInfo struct {
	Name       string       `json:"name"`
	ModifiedAt time.Time    `json:"modified_at"`
	Size       int64        `json:"size"`
	Digest     string       `json:"digest"`
	Details    ModelDetails `json:"details,omitempty"`
}

// ModelDetails represents detailed information about a model
type ModelDetails struct {
	Format            string                 `json:"format,omitempty"`
	Family            string                 `json:"family,omitempty"`
	Families          []string               `json:"families,omitempty"`
	ParameterSize     string                 `json:"parameter_size,omitempty"`
	QuantizationLevel string                 `json:"quantization_level,omitempty"`
	ParentModel       string                 `json:"parent_model,omitempty"`
	Templates         map[string]interface{} `json:"templates,omitempty"`
	System            string                 `json:"system,omitempty"`
	Options           map[string]interface{} `json:"options,omitempty"`
}

// DeleteRequest represents a request to delete a model
type DeleteRequest struct {
	Name string `json:"name" validate:"required"`
}

// DeleteResponse represents the response from a delete request
type DeleteResponse struct {
	Status string `json:"status"`
}

// ShowModelRequest represents a request to show model information
type ShowModelRequest struct {
	Name    string `json:"name" validate:"required"`
	Verbose bool   `json:"verbose,omitempty"`
}

// ShowModelResponse represents the response from a show model request
type ShowModelResponse struct {
	License    string                 `json:"license,omitempty"`
	Modelfile  string                 `json:"modelfile,omitempty"`
	Parameters string                 `json:"parameters,omitempty"`
	Template   string                 `json:"template,omitempty"`
	Details    ModelDetails           `json:"details,omitempty"`
	Options    map[string]interface{} `json:"options,omitempty"`
}

// HealthResponse represents the health check response
type HealthResponse struct {
	Status string `json:"status"`
}

// StreamChunk represents a chunk in a streaming response
type StreamChunk struct {
	Model              string    `json:"model,omitempty"`
	CreatedAt          time.Time `json:"created_at,omitempty"`
	Response           string    `json:"response,omitempty"`
	Done               bool      `json:"done"`
	Context            []int     `json:"context,omitempty"`
	TotalDuration      int64     `json:"total_duration,omitempty"`
	LoadDuration       int64     `json:"load_duration,omitempty"`
	PromptEvalCount    int       `json:"prompt_eval_count,omitempty"`
	PromptEvalDuration int64     `json:"prompt_eval_duration,omitempty"`
	EvalCount          int       `json:"eval_count,omitempty"`
	EvalDuration       int64     `json:"eval_duration,omitempty"`
	Message            *Message  `json:"message,omitempty"`
}
