package gollama

import "context"

// SessionOption configures a conversation session.
type SessionOption func(*Session)

// Session maintains multi-turn chat history for a model.
type Session struct {
	client     *Client
	model      string
	system     string
	opts       OptionSet
	messages   []Message
	maxHistory int // max user+assistant pairs; 0 = unlimited
}

// WithSessionSystem sets the session system prompt.
func WithSessionSystem(system string) SessionOption {
	return func(s *Session) { s.system = system }
}

// WithSessionOptions sets default options for session sends.
func WithSessionOptions(opts OptionSet) SessionOption {
	return func(s *Session) { s.opts = opts }
}

// WithMaxHistory limits stored turns (user+assistant pairs); oldest trimmed first.
func WithMaxHistory(pairs int) SessionOption {
	return func(s *Session) { s.maxHistory = pairs }
}

// NewSession creates a stateful chat session for the given model.
func (c *Client) NewSession(model string, opts ...SessionOption) (*Session, error) {
	if model == "" {
		model = c.defaultModel
	}
	if model == "" {
		return nil, errNoModel()
	}
	s := &Session{client: c, model: model}
	for _, opt := range opts {
		opt(s)
	}
	if s.system != "" {
		s.messages = append(s.messages, System(s.system))
	}
	return s, nil
}

// Send appends a user message, calls the model, and stores the assistant reply.
func (s *Session) Send(ctx context.Context, message string, opts ...CallOption) (*Result, error) {
	s.messages = append(s.messages, User(message))
	callOpts := append([]CallOption{WithModel(s.model), WithOptions(s.opts)}, opts...)
	result, err := s.client.ChatMessages(ctx, s.messages, callOpts...)
	if err != nil {
		// roll back user message on failure
		s.messages = s.messages[:len(s.messages)-1]
		return nil, err
	}
	s.messages = append(s.messages, Assistant(result.Text))
	s.trimHistory()
	return result, nil
}

// Messages returns a copy of the session history.
func (s *Session) Messages() []Message {
	out := make([]Message, len(s.messages))
	copy(out, s.messages)
	return out
}

// Reset clears conversation history but keeps the system prompt.
func (s *Session) Reset() {
	if s.system != "" {
		s.messages = []Message{System(s.system)}
		return
	}
	s.messages = nil
}

func (s *Session) trimHistory() {
	if s.maxHistory <= 0 {
		return
	}
	// count non-system messages as pairs
	start := 0
	if len(s.messages) > 0 && s.messages[0].Role == "system" {
		start = 1
	}
	conv := s.messages[start:]
	maxMsgs := s.maxHistory * 2
	if len(conv) <= maxMsgs {
		return
	}
	trim := len(conv) - maxMsgs
	s.messages = append(s.messages[:start], conv[trim:]...)
}