package gollama

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

	"github.com/ollama/ollama/api"
)

// OptionDoc describes a supported Ollama model option.
type OptionDoc struct {
	Key         string
	Type        string
	Description string
}

// OptionSet holds typed generation options for Ollama requests.
type OptionSet struct {
	opts api.Options
	set  map[string]struct{}
}

// NewOptions returns an empty OptionSet.
func NewOptions() OptionSet { return OptionSet{set: make(map[string]struct{})} }

// DefaultOptions returns Ollama's default option values.
func DefaultOptions() OptionSet {
	return OptionSet{opts: api.DefaultOptions(), set: make(map[string]struct{})}
}

func (o OptionSet) mark(key string) OptionSet {
	if o.set == nil {
		o.set = make(map[string]struct{})
	}
	o.set[key] = struct{}{}
	return o
}

func (o OptionSet) Temperature(v float32) OptionSet {
	o.opts.Temperature = v
	return o.mark("temperature")
}

func (o OptionSet) TopP(v float32) OptionSet {
	o.opts.TopP = v
	return o.mark("top_p")
}

func (o OptionSet) TopK(v int) OptionSet {
	o.opts.TopK = v
	return o.mark("top_k")
}

func (o OptionSet) MinP(v float32) OptionSet {
	o.opts.MinP = v
	return o.mark("min_p")
}

func (o OptionSet) TypicalP(v float32) OptionSet {
	o.opts.TypicalP = v
	return o.mark("typical_p")
}

func (o OptionSet) NumPredict(v int) OptionSet {
	o.opts.NumPredict = v
	return o.mark("num_predict")
}

func (o OptionSet) NumCtx(v int) OptionSet {
	o.opts.NumCtx = v
	return o.mark("num_ctx")
}

func (o OptionSet) NumKeep(v int) OptionSet {
	o.opts.NumKeep = v
	return o.mark("num_keep")
}

func (o OptionSet) NumBatch(v int) OptionSet {
	o.opts.NumBatch = v
	return o.mark("num_batch")
}

func (o OptionSet) NumGPU(v int) OptionSet {
	o.opts.NumGPU = v
	return o.mark("num_gpu")
}

func (o OptionSet) NumThread(v int) OptionSet {
	o.opts.NumThread = v
	return o.mark("num_thread")
}

func (o OptionSet) RepeatPenalty(v float32) OptionSet {
	o.opts.RepeatPenalty = v
	return o.mark("repeat_penalty")
}

func (o OptionSet) RepeatLastN(v int) OptionSet {
	o.opts.RepeatLastN = v
	return o.mark("repeat_last_n")
}

func (o OptionSet) PresencePenalty(v float32) OptionSet {
	o.opts.PresencePenalty = v
	return o.mark("presence_penalty")
}

func (o OptionSet) FrequencyPenalty(v float32) OptionSet {
	o.opts.FrequencyPenalty = v
	return o.mark("frequency_penalty")
}

func (o OptionSet) Seed(v int) OptionSet {
	o.opts.Seed = v
	return o.mark("seed")
}

func (o OptionSet) Stop(seq ...string) OptionSet {
	o.opts.Stop = seq
	return o.mark("stop")
}

// Merge combines base and override; fields set on override win.
func (base OptionSet) Merge(override OptionSet) OptionSet {
	out := base
	if override.set == nil {
		return out
	}
	if out.set == nil {
		out.set = make(map[string]struct{})
	}
	data, err := optionsToMap(override.opts)
	if err != nil {
		return out
	}
	for key := range override.set {
		var partial api.Options
		if err := partial.FromMap(map[string]any{key: data[key]}); err == nil {
			merged, _ := optionsToMap(out.opts)
			if v, ok := data[key]; ok {
				merged[key] = v
			}
			_ = out.opts.FromMap(merged)
			out.set[key] = struct{}{}
		}
	}
	return out
}

func (o OptionSet) toMap() (map[string]any, error) {
	return optionsToMapFiltered(o.opts, o.set)
}

func optionsToMap(opts api.Options) (map[string]any, error) {
	b, err := json.Marshal(opts)
	if err != nil {
		return nil, err
	}
	var m map[string]any
	if err := json.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	return m, nil
}

func optionsToMapFiltered(opts api.Options, set map[string]struct{}) (map[string]any, error) {
	all, err := optionsToMap(opts)
	if err != nil {
		return nil, err
	}
	if len(set) == 0 {
		return map[string]any{}, nil
	}
	out := make(map[string]any, len(set))
	for key := range set {
		if v, ok := all[key]; ok {
			out[key] = v
		}
	}
	return out, nil
}

// KnownOptions returns documentation for all supported option keys.
func KnownOptions() []OptionDoc {
	return []OptionDoc{
		{Key: "temperature", Type: "float32", Description: "Sampling temperature"},
		{Key: "top_p", Type: "float32", Description: "Nucleus sampling"},
		{Key: "top_k", Type: "int", Description: "Top-k sampling"},
		{Key: "min_p", Type: "float32", Description: "Minimum probability threshold"},
		{Key: "typical_p", Type: "float32", Description: "Typical sampling"},
		{Key: "num_predict", Type: "int", Description: "Maximum tokens to generate (-1 = unlimited)"},
		{Key: "num_ctx", Type: "int", Description: "Context window size"},
		{Key: "num_keep", Type: "int", Description: "Tokens to keep from prompt start"},
		{Key: "num_batch", Type: "int", Description: "Batch size for prompt processing"},
		{Key: "num_gpu", Type: "int", Description: "Number of GPU layers (-1 = auto)"},
		{Key: "num_thread", Type: "int", Description: "CPU threads (0 = auto)"},
		{Key: "repeat_penalty", Type: "float32", Description: "Penalty for repeated tokens"},
		{Key: "repeat_last_n", Type: "int", Description: "Lookback for repeat penalty"},
		{Key: "presence_penalty", Type: "float32", Description: "Presence penalty"},
		{Key: "frequency_penalty", Type: "float32", Description: "Frequency penalty"},
		{Key: "seed", Type: "int", Description: "Random seed (-1 = random)"},
		{Key: "stop", Type: "[]string", Description: "Stop sequences"},
	}
}

// parseModelfileParameters parses ShowResponse.Parameters into an OptionSet.
func parseModelfileParameters(parameters string) (OptionSet, error) {
	out := NewOptions()
	lines := strings.Split(parameters, "\n")
	m := make(map[string]any)
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		key, val, ok := strings.Cut(line, " ")
		if !ok {
			continue
		}
		key = strings.TrimSpace(key)
		val = strings.TrimSpace(val)
		if key == "" || val == "" {
			continue
		}
		var parsed any
		if err := json.Unmarshal([]byte(val), &parsed); err != nil {
			parsed = val
		}
		m[key] = parsed
		out.set[key] = struct{}{}
	}
	if err := out.opts.FromMap(m); err != nil {
		return OptionSet{}, fmt.Errorf("parse model parameters: %w", err)
	}
	return out, nil
}