diff --git a/src/internal/database/pics.go b/src/internal/database/pics.go index 466b43e..355e0a6 100644 --- a/src/internal/database/pics.go +++ b/src/internal/database/pics.go @@ -1,7 +1,7 @@ package database // GetLastUsedPic -func (db *Database) GetLastUsedPic() (int64, string, error) { +func (db *Database) GetLastUsedPic(increment bool) (int64, string, error) { min := db.getPicMinViewCount() var id int64 var fileName string @@ -15,7 +15,10 @@ func (db *Database) GetLastUsedPic() (int64, string, error) { row.Scan(&id, &fileName) } - db.incrementPicViewCount(id) + if increment { + // Increment the view count + db.incrementPicViewCount(id) + } return id, fileName, nil } @@ -26,14 +29,19 @@ func (db *Database) DeleteAllPics() { } // StorePic -func (db *Database) StorePic(filename string) error { +func (db *Database) StorePic(filename string) (int64, error) { min := db.getPicMinViewCount() - _, err := db.db.Exec( + exec, err := db.db.Exec( "INSERT INTO `pics` (`filename`, `viewed`) VALUES (?,?)", filename, min, ) - return err + + if err != nil { + return 0, err + } + + return exec.LastInsertId() } // DeletePic @@ -48,9 +56,9 @@ func (db *Database) ResetPics() { } // GetAllPics -func (db *Database) GetAllPics() ([]string, error) { - var pics []string - row, err := db.db.Query("SELECT filename FROM pics") +func (db *Database) GetAllPics() (map[int64]string, error) { + pics := make(map[int64]string) + row, err := db.db.Query("SELECT id, filename FROM pics") if err != nil { return pics, err } @@ -58,8 +66,9 @@ func (db *Database) GetAllPics() ([]string, error) { defer row.Close() for row.Next() { var pic string - row.Scan(&pic) - pics = append(pics, pic) + var id int64 + row.Scan(&id, &pic) + pics[id] = pic } return pics, nil diff --git a/src/internal/gomatrixbot/birds.go b/src/internal/gomatrixbot/birds.go index 86d7783..14fec59 100644 --- a/src/internal/gomatrixbot/birds.go +++ b/src/internal/gomatrixbot/birds.go @@ -19,7 +19,7 @@ const fontSize = 42 const BIRB_PATH = "/root/data/reaction_pics" func (mtrx *MtrxClient) postBirb(ctx context.Context, evt *event.Event) { - idd, birb, err := mtrx.db.GetLastUsedPic() + idd, birb, err := mtrx.db.GetLastUsedPic(true) if err != nil { mtrx.c.Log.Err(err).Msg("Failed to get birb") return @@ -58,51 +58,41 @@ func (mtrx *MtrxClient) postBirb(ctx context.Context, evt *event.Event) { }) } -func (mtrx *MtrxClient) postBirbWithText(ctx context.Context, evt *event.Event, text string, attempt int, filename string) { - idd, birb, err := mtrx.db.GetLastUsedPic() +func (mtrx *MtrxClient) postBirbWithText(ctx context.Context, evt *event.Event, text string, filename string) { + increment := true + if filename != "" { + increment = false + } + + // Do a lookup + dbId, birb, err := mtrx.db.GetLastUsedPic(increment) if err != nil { - mtrx.c.Log.Err(err).Msg("Failed to get birb") + mtrx.sendMessage(ctx, evt.RoomID, "Failed to get an image "+err.Error()+". "+birb) return } - - if filename != "" { - birb = filename - } - reader, err := os.Open(BIRB_PATH + "/" + birb) if err != nil { mtrx.sendMessage(ctx, evt.RoomID, "Failed to load image "+err.Error()+". "+birb) - mtrx.db.DeletePic(idd) - if attempt > 3 { - mtrx.db.DeletePic(idd) - return + if filename != "" { + for k, v := range mtrx.pics { + if v == birb { + dbId = k + break + } + } } - - // Retry - mtrx.postBirbWithText(ctx, evt, text, attempt+1, "") + mtrx.db.DeletePic(dbId) + delete(mtrx.pics, dbId) return } // do the swaps img, _, err := image.Decode(reader) if err != nil { - mtrx.db.DeletePic(idd) - if attempt > 3 { - mtrx.db.DeletePic(idd) - mtrx.sendMessage(ctx, evt.RoomID, "I can't read this image. I give up. "+err.Error()+". "+birb) - return - } - - // Retry - mtrx.postBirbWithText(ctx, evt, text, attempt+1, "") + mtrx.sendMessage(ctx, evt.RoomID, "I can't read this image. I give up. "+err.Error()+". "+birb) return - } - // r := img.Bounds() - // w := r.Dx() // w - // h := r.Dy() // h - imgNew := resize.Resize(1024, 1024, img, resize.Lanczos2) m := gg.NewContextForImage(imgNew) @@ -192,20 +182,31 @@ func (mtrx *MtrxClient) UpdatePics() error { } for _, file := range files { + found := false + if file.IsDir() || strings.HasPrefix(file.Name(), ".") { continue } - if contains(mtrx.pics, file.Name()) { + // check if in our map + for _, v := range mtrx.pics { + if v == file.Name() { + found = true + break + } + } + + if found { continue } - err = mtrx.db.StorePic(file.Name()) + idd, err := mtrx.db.StorePic(file.Name()) if err != nil { - mtrx.c.Log.Err(err).Msg("Failed to store pic") + return err } + mtrx.sendMessage(context.Background(), mtrx.adminRoom, "New pic added: "+file.Name()) - mtrx.pics = append(mtrx.pics, file.Name()) + mtrx.pics[idd] = file.Name() } return nil diff --git a/src/internal/gomatrixbot/commands.go b/src/internal/gomatrixbot/commands.go index 081cc09..7b4dcca 100644 --- a/src/internal/gomatrixbot/commands.go +++ b/src/internal/gomatrixbot/commands.go @@ -125,6 +125,7 @@ func (mtrx *MtrxClient) handleCommand(ctx context.Context, evt *event.Event) { return case ".reloadallpics": mtrx.db.DeleteAllPics() + mtrx.pics = make(map[int64]string, 0) mtrx.UpdatePics() mtrx.sendMessage(ctx, evt.RoomID, "Reset all pictures") return @@ -162,9 +163,9 @@ func (mtrx *MtrxClient) handleCommand(ctx context.Context, evt *event.Event) { return case ".rb": if len(args) < 2 { - mtrx.postBirbWithText(ctx, evt, mtrx.getAnyRandomQuote(true), 0, "") + mtrx.postBirbWithText(ctx, evt, mtrx.getAnyRandomQuote(true), "") } else { - mtrx.postBirbWithText(ctx, evt, mtrx.getAnyRandomQuoteSearch(true, strings.Join(args[1:], " ")), 0, "") + mtrx.postBirbWithText(ctx, evt, mtrx.getAnyRandomQuoteSearch(true, strings.Join(args[1:], " ")), "") } return case ".delquote": @@ -262,17 +263,19 @@ func (mtrx *MtrxClient) handleCommand(ctx context.Context, evt *event.Event) { filesearch := strings.ReplaceAll(args[0], ".", "") // lookup filename from quote, shuffle each time - tmpPics := make([]string, len(mtrx.pics)) - perm := rand.Perm(len(mtrx.pics)) - for i, v := range perm { - tmpPics[v] = mtrx.pics[i] + matches := make(map[int64]string) + keys := make([]int64, 0) + for k, v := range mtrx.pics { + if strings.Contains(strings.ToLower(strings.Split(v, ".")[0]), strings.ToLower(filesearch)) { + matches[k] = v + keys = append(keys, k) + } } - for _, v := range tmpPics { - if strings.Contains(strings.ToLower(strings.Split(v, ".")[0]), strings.ToLower(filesearch)) { - mtrx.postBirbWithText(ctx, evt, mtrx.getAnyRandomQuoteSearch(true, strings.Join(args[1:], " ")), 0, v) - return - } + // Shuffle key + if len(matches) > 0 { + filename := matches[keys[rand.Intn(len(keys))]] + mtrx.postBirbWithText(ctx, evt, mtrx.getAnyRandomQuoteSearch(true, strings.Join(args[1:], " ")), filename) } } } diff --git a/src/internal/gomatrixbot/matrix.go b/src/internal/gomatrixbot/matrix.go index f34f406..0dc21a9 100644 --- a/src/internal/gomatrixbot/matrix.go +++ b/src/internal/gomatrixbot/matrix.go @@ -44,7 +44,7 @@ type MtrxClient struct { savedFiles map[id.EventID]string events map[id.EventID]string recentMessages map[id.RoomID]map[int64]string - pics []string + pics map[int64]string // This is gonna be a mem leak... To fix I guess quoteCache map[id.EventID]*Quote @@ -64,7 +64,7 @@ func Run(db *database.Database) { pics, err := mtrx.db.GetAllPics() if err != nil { fmt.Println(err) - mtrx.pics = make([]string, 0) + mtrx.pics = make(map[int64]string, 0) } else { mtrx.pics = pics }