package main

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

    "github.com/mattermost/mattermost/server/public/model"
)

// ConversationStore manages conversation context in the Mattermost KV store.
type ConversationStore struct {
    api pluginKV
    mu  sync.RWMutex
}

// pluginKV is the subset of the plugin API needed for KV operations.
type pluginKV interface {
    KVGet(key string) ([]byte, *model.AppError)
    KVSet(key string, value []byte) *model.AppError
    KVDelete(key string) *model.AppError
}

// NewConversationStore creates a new store backed by the plugin API.
func NewConversationStore(api pluginKV) *ConversationStore {
    return &ConversationStore{api: api}
}

// conversationSession holds the messages for a single user+channel session.
type conversationSession struct {
    UserID    string        `json:"user_id"`
    ChannelID string        `json:"channel_id"`
    Messages  []chatMsg     `json:"messages"`
    UpdatedAt time.Time     `json:"updated_at"`
}

type chatMsg struct {
    Role    string `json:"role"`
    Content string `json:"content"`
}

const (
    maxTurns   = 10
    sessionTTL = 1 * time.Hour
)

func sessionKey(userID, channelID string) string {
    return fmt.Sprintf("ctx_%s_%s", userID, channelID)
}

// BuildMessages retrieves session history and appends the user prompt,
// returning the full message array for the Ollama API.
func (s *ConversationStore) BuildMessages(userID, channelID, prompt string) ([]ChatMessage, error) {
    s.mu.Lock()
    defer s.mu.Unlock()

    session := s.loadSession(userID, channelID)
    if session == nil {
        session = &conversationSession{
            UserID:    userID,
            ChannelID: channelID,
            Messages:  []chatMsg{},
        }
    }

    session.Messages = append(session.Messages, chatMsg{Role: "user", Content: prompt})

    // Trim to max turns (keep newest).
    if len(session.Messages) > maxTurns {
        session.Messages = session.Messages[len(session.Messages)-maxTurns:]
    }

    session.UpdatedAt = time.Now()
    _ = s.saveSession(session) // non-fatal

    result := make([]ChatMessage, 0, len(session.Messages))
    for _, m := range session.Messages {
        result = append(result, ChatMessage{Role: m.Role, Content: m.Content})
    }
    return result, nil
}

// Reset clears the session for a user+channel.
func (s *ConversationStore) Reset(userID, channelID string) error {
    s.mu.Lock()
    defer s.mu.Unlock()
    return s.api.KVDelete(sessionKey(userID, channelID))
}

func (s *ConversationStore) loadSession(userID, channelID string) *conversationSession {
    data, err := s.api.KVGet(sessionKey(userID, channelID))
    if err != nil || data == nil {
        return nil
    }
    var session conversationSession
    if err := json.Unmarshal(data, &session); err != nil {
        return nil
    }
    return &session
}

func (s *ConversationStore) saveSession(session *conversationSession) error {
    data, err := json.Marshal(session)
    if err != nil {
        return fmt.Errorf("marshal session: %w", err)
    }
    err = s.api.KVSet(sessionKey(session.UserID, session.ChannelID), data)
    return err
}

// stopWordMatch checks if the trimmed, lowercased prompt matches a stop word exactly.
func stopWordMatch(prompt string, stopWords map[string]bool) bool {
    if stopWords == nil {
        return false
    }
    return stopWords[strings.TrimSpace(strings.ToLower(prompt))]
}