package config

import (
	"errors"
	"fmt"
	"runtime"
)

const (
	MODEL_CONFIG_DEFAULT_TTL = -1
)

var validModalities = map[string]struct{}{
	"text":  {},
	"audio": {},
	"image": {},
	// 3d is the mesh modality: TRELLIS.2 takes an image and returns a GLB.
	"3d": {},
	// video is the moving-picture modality: a video DiT takes a prompt (and
	// optionally a first/last frame) and returns an encoded container. A model
	// that also scores its clip declares audio alongside it.
	"video": {},
}

// ModelCapConfig defines what modalities and features a model supports.
// Used in /v1/models to inform clients. An empty block (all zero values) is
// treated as not configured.
type ModelCapConfig struct {
	In           []string `yaml:"in"`
	Out          []string `yaml:"out"`
	Tools        bool     `yaml:"tools"`
	Reranker     bool     `yaml:"reranker"`
	Embedding    bool     `yaml:"embedding"`
	Segmentation bool     `yaml:"segmentation"`
	// Decision marks a llama.cpp decision model: it answers typed questions on
	// /v1/systemone and cannot chat. DecisionType is its "<arch>.decision.type"
	// (kev, laya, openjev, ...), which decides per-model limits upstream such as
	// the number of choice options, so clients get to see it.
	Decision     bool   `yaml:"decision"`
	DecisionType string `yaml:"decisionType"`
	// VoiceClone marks a speech model that can register NEW voices from a
	// reference clip (qwentts.cpp). It is not implied by out:[audio] — a
	// fixed-voice-pack engine like TTS.cpp/Kokoro serves speech but has no
	// endpoint to clone with, and offering the UI for it produces a 404 at best.
	VoiceClone bool `yaml:"voiceClone"`
	// Music marks an audio.cpp music generator (ACE-Step, MiniMax Music): it
	// answers /v1/music/generations, never /v1/audio/speech, so it must not
	// borrow the speech capability a text->audio In/Out pair would imply.
	// MusicEdit adds /v1/music/edits (repaint / complete / cover), which only
	// some generators implement. Separation marks a stem splitter (HTDemucs)
	// answering /v1/audio/separations.
	Music      bool `yaml:"music"`
	MusicEdit  bool `yaml:"musicEdit"`
	Separation bool `yaml:"separation"`
	// MusicScore names the score notation a music generator can be conditioned
	// on ("abc" for YuE2); empty = none. Sent as the generation route's typed
	// `score` field.
	MusicScore string `yaml:"musicScore"`
	// MusicTranscribe names the score notation a transcriber writes from a
	// recording ("abc" for SheetSage2); empty = not a transcriber. Answered on
	// /v1/music/transcriptions.
	MusicTranscribe string `yaml:"musicTranscribe"`
	// ReasoningEffort lists the reasoning-effort levels this model's chat
	// template accepts (Qwen 3.8: xhigh/medium/low), lowest-cost value LAST as
	// the template declares them. Non-empty is the signal that an incoming
	// OpenAI `reasoning_effort` field may be translated into a
	// chat_template_kwargs entry — a template that doesn't understand the value
	// raises, which surfaces as a 500, so an unlisted model is never translated
	// for. Advertised in /v1/models so a client can discover the ladder.
	ReasoningEffort []string `yaml:"reasoningEffort"`
	Context         int      `yaml:"context"`
}

// Empty returns true when all fields are at their zero values.
func (c ModelCapConfig) Empty() bool {
	return len(c.In) == 0 && len(c.Out) == 0 && !c.Tools && !c.Reranker && !c.Embedding && !c.Segmentation && !c.Decision && c.DecisionType == "" && !c.VoiceClone && !c.Music && !c.MusicEdit && !c.Separation && c.MusicScore == "" && c.MusicTranscribe == "" && len(c.ReasoningEffort) == 0 && c.Context == 0
}

// Validate checks that all modality values are recognized and context is
// non-negative. Returns an error if any value is invalid.
func (c ModelCapConfig) Validate() error {
	for _, m := range c.In {
		if _, ok := validModalities[m]; !ok {
			return fmt.Errorf("capabilities.in: invalid modality %q, must be one of: text, audio, image, 3d, video", m)
		}
	}
	for _, m := range c.Out {
		if _, ok := validModalities[m]; !ok {
			return fmt.Errorf("capabilities.out: invalid modality %q, must be one of: text, audio, image, 3d, video", m)
		}
	}
	if c.Context < 0 {
		return errors.New("capabilities.context: must be >= 0")
	}
	return nil
}

// TimeoutsConfig holds timeout settings for proxy connections
// 0 = no timeout
type TimeoutsConfig struct {
	Connect        int `yaml:"connect"`
	KeepAlive      int `yaml:"keepalive"`
	ResponseHeader int `yaml:"responseHeader"`
	TLSHandshake   int `yaml:"tlsHandshake"`
	ExpectContinue int `yaml:"expectContinue"`
	IdleConn       int `yaml:"idleConn"`
}

type ModelConfig struct {
	Cmd           string   `yaml:"cmd"`
	CmdStop       string   `yaml:"cmdStop"`
	Proxy         string   `yaml:"proxy"`
	Aliases       []string `yaml:"aliases"`
	Env           []string `yaml:"env"`
	CheckEndpoint string   `yaml:"checkEndpoint"`
	UnloadAfter   int      `yaml:"ttl"`
	Unlisted      bool     `yaml:"unlisted"`
	UseModelName  string   `yaml:"useModelName"`

	// MissingComponents names the diffusion components (vae, llm, t5xxl, ...)
	// autogen could not find on disk for this model. Pure data: the process
	// layer only reports it, because sd-server dies within seconds without
	// them and its own log is the only other place that says why.
	MissingComponents []string `yaml:"missingComponents"`

	// #179 for /v1/models
	Name        string `yaml:"name"`
	Description string `yaml:"description"`

	// Limit concurrency of HTTP requests to process
	ConcurrencyLimit int `yaml:"concurrencyLimit"`

	// Model filters see issue #174
	Filters ModelFilters `yaml:"filters"`

	// Macros: see #264
	// Model level macros take precedence over the global macros
	Macros MacroList `yaml:"macros"`

	// Metadata: see #264
	// Arbitrary metadata that can be exposed through the API
	Metadata map[string]any `yaml:"metadata"`

	// override global setting
	SendLoadingState *bool `yaml:"sendLoadingState"`

	// Timeout settings for proxy connections
	Timeouts TimeoutsConfig `yaml:"timeouts"`

	// Capabilities defines what modalities and features the model supports.
	Capabilities ModelCapConfig `yaml:"capabilities"`

	// EstVramGB is this model's predicted VRAM footprint when loaded (weights +
	// KV reserve + compute buffer + overhead), as computed by the autogen sizer
	// at generate time. It is the admission input for VRAM-budget-aware
	// multi-load (see Config.VramBudgetGB): the router sums it over the resident
	// set to decide whether a new model fits alongside them or something must be
	// evicted first. 0 = unknown (hand-written config, or a class the sizer does
	// not model), which the router treats conservatively — see
	// internal/router/group.go.
	EstVramGB float64 `yaml:"estVramGB"`

	// EstRamGB is the sizer's companion figure: the share of the weights (plus
	// KV, when it is kept in system memory) that does NOT fit on the GPU and is
	// served from RAM. 0 for a fully-offloaded model. The router charges it
	// against Config.RamBudgetGB (every resident, whichever pool) and evicts
	// ram-pool models to make it fit; the dashboard shows it per model so the
	// cost of a partial offload is visible before loading.
	EstRamGB float64 `yaml:"estRamGB"`

	// SlotCachePreamble opts this model out of PREAMBLE caches (the shared
	// system+tools KV seed that is minted once per agent and reused by every
	// load) while leaving its per-conversation snapshots alone. nil/absent =>
	// on, so a model that opted into slot persistence at all gets both halves.
	// The whole-feature per-model gate stays the cmd's --slot-save-path; this is
	// the finer knob, because a preamble cache is minted unprompted (a synthetic
	// prefill) and is exempt from the LRU disk cap.
	SlotCachePreamble *bool `yaml:"slotCachePreamble"`

	// AudioCpp carries the one thing audiocpp_server cannot be told on its
	// command line. It has no --model flag at all: the model is named only in a
	// JSON config file passed as --config, while every other knob is an argv
	// flag that overrides the file. Keeping the entry here instead of writing a
	// JSON file per model at generate time means one source of truth, no
	// generated files to keep in step with the config, and nothing left behind
	// when a model is renamed or removed. The spawn hook materializes it (see
	// internal/server/audiocppconfig.go).
	AudioCpp AudioCppConfig `yaml:"audiocpp"`

	// PromptEnhancer names the Config.PromptEnhancers entry this IMAGE model
	// sends its prompt through before rendering. Empty (every non-image model,
	// and every image model that has not opted in) => none.
	//
	// A dangling id is not an error: an enhancer whose model was deleted resolves
	// to nothing and the model renders the prompt as typed. Failing config load
	// over a rewrite convenience would be the wrong trade.
	PromptEnhancer string `yaml:"promptEnhancer"`

	// PromptEnhancerEdit is the enhancer used when the request carries a
	// reference image, replacing PromptEnhancer for that direction only. Empty =>
	// PromptEnhancer covers both. See autogen.Override.PromptEnhancerEdit.
	PromptEnhancerEdit string `yaml:"promptEnhancerEdit"`

	// PromptEnhancerPrompt / PromptEnhancerEditPrompt are THIS model's system
	// prompts for the two directions, overriding the shared
	// Config.PromptEnhancers entry's. Empty => use the shared one. See
	// autogen.Override.PromptEnhancerPrompt for why the prompt belongs to the
	// image model rather than to the rewriter.
	PromptEnhancerPrompt     string `yaml:"promptEnhancerPrompt"`
	PromptEnhancerEditPrompt string `yaml:"promptEnhancerEditPrompt"`

	// Copy of HealthCheckTimeout from global config
	HealthCheckTimeout int `yaml:"healthCheckTimeout"`
}

// AudioCppConfig is ONE entry of audio.cpp's models[] array: the fields its
// server requires (id comes from the model id) plus the task that selects which
// route the model answers on. Deliberately not a passthrough for the rest of
// that schema - lazy loading, the residency cap, idle unloads and the memory
// guard are audio.cpp's own scheduler, which is contained rather than
// configured here (see internal/autogen/audiocpp.go).
type AudioCppConfig struct {
	// Family is audio.cpp's model family id (moss_tts_nano, qwen3_asr, ...),
	// read from the gguf's own audiocpp.model_spec.family header key.
	Family string `yaml:"family"`
	// Path is the gguf (or model directory) the server loads.
	Path string `yaml:"path"`
	// Task selects which surface this process serves: "tts", "asr", "gen"
	// (music generation), "sep" (source separation) or "midi" (score
	// transcription). Blank lets audio.cpp
	// apply its own default ("tts").
	Task string `yaml:"task"`
	// SessionOptions are audio.cpp's per-model session_options, passed through
	// verbatim. A directory model whose package holds several quants
	// (MiniMax Music 3) is told which component files to load here, since its
	// built-in defaults name a quant that may not be on disk.
	SessionOptions map[string]string `yaml:"sessionOptions"`
}

// Empty reports whether this model carries no audio.cpp block at all. Family
// and Path are both required by the server, so a half-filled block is as
// unusable as an absent one and is treated the same way.
func (a AudioCppConfig) Empty() bool {
	return a.Family == "" || a.Path == ""
}

func (m *ModelConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
	type rawModelConfig ModelConfig
	defaults := rawModelConfig{
		Cmd:              "",
		CmdStop:          "",
		Proxy:            "http://localhost:${PORT}",
		Aliases:          []string{},
		Env:              []string{},
		CheckEndpoint:    "/health",
		UnloadAfter:      MODEL_CONFIG_DEFAULT_TTL, // use GlobalTTL
		Unlisted:         false,
		UseModelName:     "",
		ConcurrencyLimit: 0,
		Name:             "",
		Description:      "",

		// matches http.DefaultTransport
		Timeouts: TimeoutsConfig{
			Connect:        30,
			KeepAlive:      30,
			ResponseHeader: 0,
			TLSHandshake:   10,
			ExpectContinue: 1,
			IdleConn:       90,
		},
	}

	// the default cmdStop to taskkill /f /t /pid ${PID}
	if runtime.GOOS == "windows" {
		defaults.CmdStop = "taskkill /f /t /pid ${PID}"
	}

	if err := unmarshal(&defaults); err != nil {
		return err
	}

	*m = ModelConfig(defaults)
	return nil
}

func (m *ModelConfig) SanitizedCommand() ([]string, error) {
	return SanitizeCommand(m.Cmd)
}

// ModelFilters embeds Filters and adds legacy support for strip_params field
// See issue #174
type ModelFilters struct {
	Filters `yaml:",inline"`
}

func (m *ModelFilters) UnmarshalYAML(unmarshal func(interface{}) error) error {
	type rawModelFilters ModelFilters
	defaults := rawModelFilters{}

	if err := unmarshal(&defaults); err != nil {
		return err
	}

	// Try to unmarshal with the old field name for backwards compatibility
	if defaults.StripParams == "" {
		var legacy struct {
			StripParams string `yaml:"strip_params"`
		}
		if legacyErr := unmarshal(&legacy); legacyErr != nil {
			return errors.New("failed to unmarshal legacy filters.strip_params: " + legacyErr.Error())
		}
		defaults.StripParams = legacy.StripParams
	}

	*m = ModelFilters(defaults)
	return nil
}

// SanitizedStripParams wraps Filters.SanitizedStripParams for backwards compatibility
// Returns ([]string, error) to match existing API
func (f ModelFilters) SanitizedStripParams() ([]string, error) {
	return f.Filters.SanitizedStripParams(), nil
}