package main

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

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

const kvPrefix = "ls_session_"

// SessionData represents a tracked agent session.
type SessionData struct {
	SessionKey  string        `json:"session_key"`
	PostID      string        `json:"post_id"`
	ChannelID   string        `json:"channel_id"`
	RootID      string        `json:"root_id,omitempty"`
	AgentID     string        `json:"agent_id"`
	Status      string        `json:"status"` // active, done, error, interrupted
	Lines       []string      `json:"lines"`
	ElapsedMs   int64         `json:"elapsed_ms"`
	TokenCount  int           `json:"token_count"`
	Children    []SessionData `json:"children,omitempty"`
	StartTimeMs int64         `json:"start_time_ms"`
}

// Store wraps Mattermost KV store operations for session persistence.
type Store struct {
	api plugin.API
}

// NewStore creates a new Store instance.
func NewStore(api plugin.API) *Store {
	return &Store{api: api}
}

// encodeKey URL-encodes a session key for safe KV storage.
func encodeKey(sessionKey string) string {
	return kvPrefix + url.PathEscape(sessionKey)
}

// SaveSession stores a session in the KV store.
func (s *Store) SaveSession(sessionKey string, data SessionData) error {
	b, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("marshal session: %w", err)
	}
	appErr := s.api.KVSet(encodeKey(sessionKey), b)
	if appErr != nil {
		return fmt.Errorf("kv set: %s", appErr.Error())
	}
	return nil
}

// GetSession retrieves a session from the KV store.
func (s *Store) GetSession(sessionKey string) (*SessionData, error) {
	b, appErr := s.api.KVGet(encodeKey(sessionKey))
	if appErr != nil {
		return nil, fmt.Errorf("kv get: %s", appErr.Error())
	}
	if b == nil {
		return nil, nil
	}
	var data SessionData
	if err := json.Unmarshal(b, &data); err != nil {
		return nil, fmt.Errorf("unmarshal session: %w", err)
	}
	return &data, nil
}

// DeleteSession removes a session from the KV store.
func (s *Store) DeleteSession(sessionKey string) error {
	appErr := s.api.KVDelete(encodeKey(sessionKey))
	if appErr != nil {
		return fmt.Errorf("kv delete: %s", appErr.Error())
	}
	return nil
}

// ListActiveSessions returns all active sessions from the KV store.
func (s *Store) ListActiveSessions() ([]SessionData, error) {
	var sessions []SessionData
	page := 0
	perPage := 100

	for {
		keys, appErr := s.api.KVList(page, perPage)
		if appErr != nil {
			return nil, fmt.Errorf("kv list: %s", appErr.Error())
		}
		if len(keys) == 0 {
			break
		}

		for _, key := range keys {
			if !strings.HasPrefix(key, kvPrefix) {
				continue
			}
			b, getErr := s.api.KVGet(key)
			if getErr != nil || b == nil {
				continue
			}
			var data SessionData
			if err := json.Unmarshal(b, &data); err != nil {
				continue
			}
			if data.Status == "active" {
				sessions = append(sessions, data)
			}
		}
		page++
	}

	return sessions, nil
}