package memorytool

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"strings"
	"time"
)

var ErrInvalidToken = errors.New("invalid memory tool token")

type Claims struct {
	ConversationID string `json:"conversation_id"`
	TaskRef        string `json:"task_ref"`
	ExpiresAt      int64  `json:"expires_at"`
}

type TokenSigner struct {
	secret []byte
	now    func() time.Time
}

func NewTokenSigner(secret []byte) (*TokenSigner, error) {
	if len(secret) < 32 {
		return nil, errors.New("memory tool signing secret must contain at least 32 bytes")
	}
	return &TokenSigner{secret: append([]byte(nil), secret...), now: time.Now}, nil
}

func (s *TokenSigner) Issue(conversationID, taskRef string, expiresAt time.Time) (string, error) {
	if strings.TrimSpace(conversationID) == "" || strings.TrimSpace(taskRef) == "" || !expiresAt.After(s.now()) {
		return "", ErrInvalidToken
	}
	payload, err := json.Marshal(Claims{ConversationID: conversationID, TaskRef: taskRef, ExpiresAt: expiresAt.UTC().Unix()})
	if err != nil {
		return "", err
	}
	encoded := base64.RawURLEncoding.EncodeToString(payload)
	return "v1." + encoded + "." + s.signature(encoded), nil
}

func (s *TokenSigner) Verify(token string) (Claims, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 || parts[0] != "v1" {
		return Claims{}, ErrInvalidToken
	}
	expected := s.signature(parts[1])
	if !hmac.Equal([]byte(expected), []byte(parts[2])) {
		return Claims{}, ErrInvalidToken
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return Claims{}, ErrInvalidToken
	}
	var claims Claims
	if err := json.Unmarshal(payload, &claims); err != nil || claims.ConversationID == "" || claims.TaskRef == "" || claims.ExpiresAt <= s.now().UTC().Unix() {
		return Claims{}, ErrInvalidToken
	}
	return claims, nil
}

func (s *TokenSigner) signature(payload string) string {
	mac := hmac.New(sha256.New, s.secret)
	_, _ = mac.Write([]byte("memory-tool-token-v1." + payload))
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}