package playgroundapp

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/ardanlabs/kronk/sdk/kronk/model"
)

// =============================================================================

// SessionRequest represents the request to create a playground session.
type SessionRequest struct {
	ModelID        string        `json:"model_id"`
	TemplateMode   string        `json:"template_mode"`
	TemplateName   string        `json:"template_name"`
	TemplateScript string        `json:"template_script"`
	Config         SessionConfig `json:"config"`
}

// Decode implements the decoder interface.
func (s *SessionRequest) Decode(data []byte) error {
	return json.Unmarshal(data, s)
}

// Validate checks the request.
func (s *SessionRequest) Validate() error {
	if s.ModelID == "" {
		return errors.New("model_id is required")
	}

	switch s.TemplateMode {
	case "builtin", "custom":
	case "":
		s.TemplateMode = "builtin"
	default:
		return fmt.Errorf("invalid template_mode: %s", s.TemplateMode)
	}

	if s.TemplateMode == "custom" && s.TemplateScript == "" {
		return errors.New("template_script is required when template_mode is custom")
	}

	if s.TemplateName != "" {
		if err := validateTemplateName(s.TemplateName); err != nil {
			return err
		}
	}

	return s.Config.Validate()
}

// SessionConfig represents model configuration overrides. Pointer fields allow
// distinguishing "not set by user" (nil) from an explicit value, so that only
// user-provided overrides are merged on top of the catalog-resolved base config.
type SessionConfig struct {
	ContextWindow       *int                      `json:"context_window"`
	NBatch              *int                      `json:"nbatch"`
	NUBatch             *int                      `json:"nubatch"`
	NSeqMax             *int                      `json:"nseq_max"`
	FlashAttention      *model.FlashAttentionType `json:"flash_attention"`
	CacheTypeK          *model.GGMLType           `json:"cache_type_k"`
	CacheTypeV          *model.GGMLType           `json:"cache_type_v"`
	NGpuLayers          *int                      `json:"ngpu_layers"`
	IncrementalCache    *bool                     `json:"incremental_cache"`
	RopeScaling         *model.RopeScalingType    `json:"rope_scaling_type"`
	RopeFreqBase        *float32                  `json:"rope_freq_base"`
	RopeFreqScale       *float32                  `json:"rope_freq_scale"`
	YarnExtFactor       *float32                  `json:"yarn_ext_factor"`
	YarnAttnFactor      *float32                  `json:"yarn_attn_factor"`
	YarnBetaFast        *float32                  `json:"yarn_beta_fast"`
	YarnBetaSlow        *float32                  `json:"yarn_beta_slow"`
	YarnOrigCtx         *int                      `json:"yarn_orig_ctx"`
	SplitMode           *model.SplitMode          `json:"split_mode"`
	Devices             []string                  `json:"devices"`
	MainGPU             *int                      `json:"main_gpu"`
	TensorSplit         []float32                 `json:"tensor_split"`
	OpOffloadMinBatch   *int                      `json:"op_offload_min_batch"`
	TensorBuftOverrides []string                  `json:"tensor_buft_overrides"`
	DraftModelID        *string                   `json:"draft_model_id"`
	DraftNDraft         *int                      `json:"draft_ndraft"`
}

// ApplyTo merges user overrides onto a base model config. Only fields
// explicitly provided by the user (non-nil pointers) are applied.
func (sc SessionConfig) ApplyTo(cfg model.Config) model.Config {
	if sc.ContextWindow != nil {
		cfg.PtrContextWindow = sc.ContextWindow
	}
	if sc.NBatch != nil {
		cfg.PtrNBatch = sc.NBatch
	}
	if sc.NUBatch != nil {
		cfg.PtrNUBatch = sc.NUBatch
	}
	if sc.NSeqMax != nil {
		cfg.PtrNSeqMax = sc.NSeqMax
	}
	if sc.FlashAttention != nil {
		cfg.FlashAttention = *sc.FlashAttention
	}
	if sc.CacheTypeK != nil {
		cfg.CacheTypeK = *sc.CacheTypeK
	}
	if sc.CacheTypeV != nil {
		cfg.CacheTypeV = *sc.CacheTypeV
	}
	if sc.NGpuLayers != nil {
		cfg.PtrNGpuLayers = sc.NGpuLayers
	}
	if sc.IncrementalCache != nil {
		cfg.PtrIncrementalCache = sc.IncrementalCache
	}
	if sc.RopeScaling != nil {
		cfg.RopeScaling = *sc.RopeScaling
	}
	if sc.RopeFreqBase != nil {
		cfg.PtrRopeFreqBase = sc.RopeFreqBase
	}
	if sc.RopeFreqScale != nil {
		cfg.PtrRopeFreqScale = sc.RopeFreqScale
	}
	if sc.YarnExtFactor != nil {
		cfg.PtrYarnExtFactor = sc.YarnExtFactor
	}
	if sc.YarnAttnFactor != nil {
		cfg.PtrYarnAttnFactor = sc.YarnAttnFactor
	}
	if sc.YarnBetaFast != nil {
		cfg.PtrYarnBetaFast = sc.YarnBetaFast
	}
	if sc.YarnBetaSlow != nil {
		cfg.PtrYarnBetaSlow = sc.YarnBetaSlow
	}
	if sc.YarnOrigCtx != nil {
		cfg.PtrYarnOrigCtx = sc.YarnOrigCtx
	}
	if sc.SplitMode != nil {
		cfg.PtrSplitMode = sc.SplitMode
	}
	if len(sc.Devices) > 0 {
		cfg.Devices = sc.Devices
	}
	if sc.MainGPU != nil {
		cfg.PtrMainGPU = sc.MainGPU
	}
	if len(sc.TensorSplit) > 0 {
		cfg.TensorSplit = sc.TensorSplit
	}
	if sc.OpOffloadMinBatch != nil {
		cfg.PtrOpOffloadMinBatch = sc.OpOffloadMinBatch
	}
	if len(sc.TensorBuftOverrides) > 0 {
		cfg.TensorBuftOverrides = sc.TensorBuftOverrides
	}
	// Draft model overrides. Three shapes are supported:
	//   - draft_model_id set & non-empty → separate-GGUF drafter (file
	//     paths are resolved in the handler, which also forces nseq-max=1).
	//   - draft_model_id set & empty     → clear any drafter.
	//   - draft_ndraft only (no id)      → MTP nDraft override; tunes the
	//     starting draft-token count for the target's auto-detected MTP
	//     head without supplying a separate draft GGUF.
	switch {
	case sc.DraftModelID != nil && *sc.DraftModelID == "":
		cfg.DraftModel = nil
	case sc.DraftModelID != nil || sc.DraftNDraft != nil:
		if cfg.DraftModel == nil {
			cfg.DraftModel = &model.DraftModelConfig{}
		}
	}
	if sc.DraftNDraft != nil && cfg.DraftModel != nil {
		cfg.DraftModel.NDraft = *sc.DraftNDraft
	}
	return cfg
}

// HasOverrides reports whether any configuration field was explicitly
// provided by the user.
func (sc SessionConfig) HasOverrides() bool {
	return sc.ContextWindow != nil ||
		sc.NBatch != nil ||
		sc.NUBatch != nil ||
		sc.NSeqMax != nil ||
		sc.FlashAttention != nil ||
		sc.CacheTypeK != nil ||
		sc.CacheTypeV != nil ||
		sc.NGpuLayers != nil ||
		sc.IncrementalCache != nil ||
		sc.RopeScaling != nil ||
		sc.RopeFreqBase != nil ||
		sc.RopeFreqScale != nil ||
		sc.YarnExtFactor != nil ||
		sc.YarnAttnFactor != nil ||
		sc.YarnBetaFast != nil ||
		sc.YarnBetaSlow != nil ||
		sc.YarnOrigCtx != nil ||
		sc.SplitMode != nil ||
		sc.Devices != nil ||
		sc.MainGPU != nil ||
		sc.TensorSplit != nil ||
		sc.OpOffloadMinBatch != nil ||
		sc.DraftModelID != nil ||
		sc.DraftNDraft != nil ||
		sc.TensorBuftOverrides != nil
}

// HasOverrides reports whether the request contains any config or template
// overrides that require a separate model instance from the Chat path.
func (s SessionRequest) HasOverrides() bool {
	return s.Config.HasOverrides() ||
		s.TemplateMode == "custom" ||
		s.TemplateName != "" ||
		s.TemplateScript != ""
}

// Validate checks the configuration bounds.
func (sc SessionConfig) Validate() error {
	if sc.ContextWindow != nil && (*sc.ContextWindow < 1 || *sc.ContextWindow > 131072) {
		return fmt.Errorf("context-window must be between 1 and 131072, got %d", *sc.ContextWindow)
	}

	if sc.NBatch != nil && (*sc.NBatch < 1 || *sc.NBatch > 16384) {
		return fmt.Errorf("nbatch must be between 1 and 16384, got %d", *sc.NBatch)
	}

	if sc.NUBatch != nil && (*sc.NUBatch < 1 || *sc.NUBatch > 16384) {
		return fmt.Errorf("nubatch must be between 1 and 16384, got %d", *sc.NUBatch)
	}

	if sc.NSeqMax != nil && (*sc.NSeqMax < 1 || *sc.NSeqMax > 64) {
		return fmt.Errorf("nseq-max must be between 1 and 64, got %d", *sc.NSeqMax)
	}

	if sc.OpOffloadMinBatch != nil && *sc.OpOffloadMinBatch < 0 {
		return fmt.Errorf("op-offload-min-batch must be >= 0, got %d", *sc.OpOffloadMinBatch)
	}

	// draft_ndraft is valid both with a separate draft model (draft_model_id)
	// and on its own as an MTP nDraft override on the target's auto-detected
	// MTP head.
	if sc.DraftNDraft != nil {
		if *sc.DraftNDraft < 1 || *sc.DraftNDraft > 20 {
			return fmt.Errorf("draft_ndraft must be between 1 and 20, got %d", *sc.DraftNDraft)
		}
	}

	return nil
}

// =============================================================================

// SessionResponse represents the response from creating a session.
type SessionResponse struct {
	SessionID       string         `json:"session_id"`
	CacheKey        string         `json:"cache_key,omitempty"`
	Status          string         `json:"status"`
	EffectiveConfig map[string]any `json:"effective_config"`
}

// Encode implements the encoder interface.
func (s SessionResponse) Encode() ([]byte, string, error) {
	data, err := json.Marshal(s)
	return data, "application/json", err
}

// =============================================================================

// SessionDeleteResponse represents the response from deleting a session.
type SessionDeleteResponse struct {
	Status string `json:"status"`
}

// Encode implements the encoder interface.
func (s SessionDeleteResponse) Encode() ([]byte, string, error) {
	data, err := json.Marshal(s)
	return data, "application/json", err
}

// =============================================================================

func validateTemplateName(name string) error {
	if name == "" {
		return fmt.Errorf("missing template name")
	}
	if len(name) > 255 {
		return fmt.Errorf("template name too long")
	}
	if strings.Contains(name, "..") || strings.Contains(name, "/") || strings.Contains(name, "\\") {
		return fmt.Errorf("invalid template name: %s", name)
	}
	if name[0] == '.' {
		return fmt.Errorf("template name must not start with a dot")
	}
	return nil
}

// =============================================================================

// ChatRequest represents a playground chat request.
type ChatRequest struct {
	SessionID string `json:"session_id"`
}

// Decode implements the decoder interface.
func (c *ChatRequest) Decode(data []byte) error {
	return json.Unmarshal(data, c)
}
