| 142 | }; |
| 143 | |
| 144 | export const MessageBranch = ({ |
| 145 | defaultBranch = 0, |
| 146 | onBranchChange, |
| 147 | className, |
| 148 | ...props |
| 149 | }: MessageBranchProps) => { |
| 150 | const [currentBranch, setCurrentBranch] = useState(defaultBranch); |
| 151 | const [branches, setBranches] = useState<ReactElement[]>([]); |
| 152 | |
| 153 | const handleBranchChange = useCallback( |
| 154 | (newBranch: number) => { |
| 155 | setCurrentBranch(newBranch); |
| 156 | onBranchChange?.(newBranch); |
| 157 | }, |
| 158 | [onBranchChange], |
| 159 | ); |
| 160 | |
| 161 | const goToPrevious = useCallback(() => { |
| 162 | const newBranch = |
| 163 | currentBranch > 0 ? currentBranch - 1 : branches.length - 1; |
| 164 | handleBranchChange(newBranch); |
| 165 | }, [currentBranch, branches.length, handleBranchChange]); |
| 166 | |
| 167 | const goToNext = useCallback(() => { |
| 168 | const newBranch = |
| 169 | currentBranch < branches.length - 1 ? currentBranch + 1 : 0; |
| 170 | handleBranchChange(newBranch); |
| 171 | }, [currentBranch, branches.length, handleBranchChange]); |
| 172 | |
| 173 | const contextValue = useMemo<MessageBranchContextType>( |
| 174 | () => ({ |
| 175 | branches, |
| 176 | currentBranch, |
| 177 | goToNext, |
| 178 | goToPrevious, |
| 179 | setBranches, |
| 180 | totalBranches: branches.length, |
| 181 | }), |
| 182 | [branches, currentBranch, goToNext, goToPrevious], |
| 183 | ); |
| 184 | |
| 185 | return ( |
| 186 | <MessageBranchContext.Provider value={contextValue}> |
| 187 | <div |
| 188 | className={cn("grid w-full gap-2 [&>div]:pb-0", className)} |
| 189 | {...props} |
| 190 | /> |
| 191 | </MessageBranchContext.Provider> |
| 192 | ); |
| 193 | }; |
| 194 | |
| 195 | export type MessageBranchContentProps = HTMLAttributes<HTMLDivElement>; |
| 196 | |