package laniakea import ( "math/rand/v2" "sync" "sync/atomic" "git.scuroneko.dev/scuroneko/laniakea/tgapi" ) type draftIDGenerator interface { Next() uint64 } // RandomDraftIDGenerator generates draft IDs using math/rand/v2. // // Suitable for general use thanks to the wide 64-bit value space. Not suitable // for security-sensitive purposes — use crypto/rand if unpredictability against // an adversary matters. type RandomDraftIDGenerator struct{} // Next returns a random 64-bit unsigned integer. func (g *RandomDraftIDGenerator) Next() uint64 { return rand.Uint64() } // LinearDraftIDGenerator generates draft IDs using a monotonically increasing counter. // Useful for debugging, persistence, or when drafts must be ordered. type LinearDraftIDGenerator struct { lastID atomic.Uint64 } // Next returns the next linear ID, atomically incremented. func (g *LinearDraftIDGenerator) Next() uint64 { return g.lastID.Add(1) } // DraftProvider manages a collection of Drafts and a shared draft ID generator. // // DraftProvider is safe for concurrent use. type DraftProvider struct { mu sync.RWMutex api *tgapi.API drafts map[uint64]*Draft generator draftIDGenerator } // NewRandomDraftProvider creates a new DraftProvider using random draft IDs. // // The provider will use random numbers for draft IDs. // All drafts created via this provider will have unpredictable, unique IDs. func NewRandomDraftProvider(api *tgapi.API) *DraftProvider { return &DraftProvider{ api: api, generator: &RandomDraftIDGenerator{}, drafts: make(map[uint64]*Draft), } } // NewLinearDraftProvider creates a new DraftProvider using linear (incrementing) draft IDs. // // startValue is the initial value for the counter. Use 0 for fresh start, or a known // value to resume from persisted state. // // This is useful when you need to store draft IDs externally (e.g., in a database) // and want to reconstruct drafts after restart. func NewLinearDraftProvider(api *tgapi.API, startValue uint64) *DraftProvider { g := &LinearDraftIDGenerator{} g.lastID.Store(startValue) return &DraftProvider{ api: api, generator: g, drafts: make(map[uint64]*Draft), } } // GetDraft retrieves a draft by its ID. // // Returns the draft and true if found, or nil and false if not found. func (p *DraftProvider) GetDraft(id uint64) (*Draft, bool) { p.mu.RLock() defer p.mu.RUnlock() draft, ok := p.drafts[id] return draft, ok } // FlushAll sends all pending drafts as final messages and clears them. // // If one or more drafts fail to send, FlushAll still attempts all drafts and // returns the first encountered error. // // After successful flush, each draft is removed from the provider and cleared. func (p *DraftProvider) FlushAll() error { p.mu.RLock() drafts := make([]*Draft, 0, len(p.drafts)) for _, draft := range p.drafts { drafts = append(drafts, draft) } p.mu.RUnlock() var firstErr error for _, draft := range drafts { if err := draft.Flush(); err != nil && firstErr == nil { firstErr = err } } return firstErr } // Draft represents a single message draft that can be edited and flushed. // // Drafts are safe to use from a single goroutine. Multiple goroutines must // synchronize access manually. // // Drafts are automatically removed from the provider's map when Flush() succeeds. type Draft struct { api *tgapi.API provider *DraftProvider chatID int64 messageThreadID int parseMode tgapi.ParseMode entities []tgapi.MessageEntity // ID uniquely identifies the draft within its provider. ID uint64 // Message contains the current draft text. Message string } // NewDraft creates a new draft with the provided parse mode. // // The caller must set a chat with SetChat before Push or Flush. func (p *DraftProvider) NewDraft(parseMode tgapi.ParseMode) *Draft { p.mu.Lock() defer p.mu.Unlock() var id uint64 for { id = p.generator.Next() if id == 0 { continue } if _, exists := p.drafts[id]; !exists { break } } draft := &Draft{ api: p.api, provider: p, parseMode: parseMode, ID: id, Message: "", } p.drafts[id] = draft return draft } // SetChat overrides the draft's target chat and message thread. // // This is useful for sending a draft to a different chat than the provider's default. func (d *Draft) SetChat(chatID int64, messageThreadID int) *Draft { d.chatID = chatID d.messageThreadID = messageThreadID return d } // SetEntities replaces the draft's message entities. // // The entities slice is copied. func (d *Draft) SetEntities(entities []tgapi.MessageEntity) *Draft { d.entities = append([]tgapi.MessageEntity(nil), entities...) return d } // Push appends text to the draft and attempts to update the server-side draft. // // Returns an error if the Telegram API rejects the update (e.g., due to network issues). // The draft's Message field is always updated, even if the API call fails. // // Use this method to build the message incrementally. func (d *Draft) Push(text string) error { return d.push(text) } // GetMessage returns the current content of the draft. // // Useful for inspection, logging, or validation before flushing. func (d *Draft) GetMessage() string { return d.Message } // Clear resets the draft's message content to empty string. // // Does not affect server-side draft — use Flush() for that. func (d *Draft) Clear() { d.Message = "" } // Delete removes the draft from its provider and clears its content. // // You may call it manually if you want to cancel a draft without sending it. func (d *Draft) Delete() { if d.provider != nil { d.provider.mu.Lock() delete(d.provider.drafts, d.ID) d.provider.mu.Unlock() } d.Clear() } // Flush sends the draft as a final message and clears it locally. // // If successful: // - The message is sent to Telegram. // - The draft's content is cleared. // - The draft is removed from the provider's map. // // If an error occurs: // - The message is NOT sent. // - The draft remains in the provider and retains its content. // - You can call Flush() again to retry. // // If the draft is empty, Flush() returns nil without calling the API. func (d *Draft) Flush() error { if d.Message == "" { d.Delete() return nil } if d.chatID == 0 { return ErrDraftChatIDZero } if err := validateMessageText(d.Message); err != nil { return err } if d.api == nil { return ErrAPIIsNil } params := tgapi.SendMessage{ ChatID: d.chatID, ParseMode: d.parseMode, Entities: d.entities, Text: d.Message, } if d.messageThreadID > 0 { params.MessageThreadID = d.messageThreadID } _, err := d.api.SendMessage(params) if err == nil { d.Delete() } return err } // The candidate Message (current content + new text) is validated before any // mutation, so a validation failure leaves the draft unchanged. After the // validation passes, Message is committed locally regardless of whether the // API call succeeds (per the Push docs: local state reflects the user's // intent, network failures can be retried). func (d *Draft) push(text string) error { if d.chatID == 0 { return ErrDraftChatIDZero } candidate := d.Message + text if err := validateMessageText(candidate); err != nil { return err } d.Message = candidate if d.api == nil { return ErrAPIIsNil } params := tgapi.SendMessageDraft{ ChatID: d.chatID, DraftID: d.ID, Text: d.Message, ParseMode: d.parseMode, Entities: d.entities, } if d.messageThreadID > 0 { params.MessageThreadID = d.messageThreadID } _, err := d.api.SendMessageDraft(params) return err }