fix chat focus

feature/chat/chat-client/src/commonMain/kotlin/ru/shadowsparky/chat/presentation/ChatScreen.kt

@@ -45,6 +45,7 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshotFlow
 import androidx.compose.runtime.withFrameNanos
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
@@ -68,6 +69,7 @@ import androidx.compose.ui.text.style.TextDecoration
 import androidx.compose.ui.text.withLink
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.dp
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 import org.jetbrains.compose.resources.painterResource
 import org.jetbrains.compose.resources.stringResource
@@ -83,11 +85,7 @@ import ru.shadowsparky.ui.components.isTv
 import ru.shadowsparky.ui.theme.focusBorder
 
 @Composable
-fun ChatContent(
+fun ChatContent(component: ChatComponent, paddingValues: PaddingValues, state: ChatState) {
-    component: ChatComponent,
-    paddingValues: PaddingValues,
-    state: ChatState
-) {
     Box(
         modifier = Modifier.fillMaxSize()
             .padding(
@@ -144,19 +142,23 @@ private fun ChatContent(
     val listState = rememberLazyListState()
     var textInput by remember { mutableStateOf("") }
     val textFieldFocusRequester = remember { FocusRequester() }
-
     LaunchedEffect(state.messages.size) {
         if (state.messages.isNotEmpty()) {
             listState.animateScrollToItem(0)
         }
     }
-
     LaunchedEffect(state.messages) { textInput = "" }
-
     Column(modifier = Modifier.fillMaxSize()) {
+        val scope = rememberCoroutineScope()
+        val focusRequesters = remember { mutableMapOf<Any, FocusRequester>() }
+        fun requesterFor(message: Message): FocusRequester {
+            val key = message.id ?: message.hashCode()
+            return focusRequesters.getOrPut(key) { FocusRequester() }
+        }
         LazyColumn(
             state = listState,
-            modifier = Modifier.weight(1f)
+            modifier = Modifier
+                .weight(1f)
                 .fillMaxWidth(),
             contentPadding = contentPadding,
             verticalArrangement = Arrangement.spacedBy(8.dp),
@@ -164,15 +166,48 @@ private fun ChatContent(
         ) {
             itemsIndexed(
                 items = state.messages,
-                key = { _, item -> item.id ?: item.hashCode() }
+                key = { _, item -> item.id ?: item.hashCode() },
             ) { index, message ->
+                val requester = requesterFor(message)
                 MessageBubble(
+                    itemFocusRequester = requester,
                     message = message,
                     onParseMessage = onParseMessage,
                     onClickLink = onClickLink,
-                    onKeyEvent = {
+                    onKeyEvent = { event ->
-                        if (index == 0) {
+                        when (event.key) {
-                            if (it.key == Key.DirectionDown) textFieldFocusRequester.requestFocus()
+                            Key.DirectionUp -> {
+                                val targetIndex = index + 1
+                                if (targetIndex < state.messages.size) {
+                                    val target = state.messages[targetIndex]
+                                    val targetRequester = requesterFor(target)
+                                    scope.launch {
+                                        listState.scrollToItem(targetIndex)
+                                        snapshotFlow {
+                                            listState.layoutInfo.visibleItemsInfo
+                                                .any { it.index == targetIndex }
+                                        }.first { it }
+                                        targetRequester.requestFocus()
+                                    }
+                                }
+                            }
+                            Key.DirectionDown -> {
+                                val targetIndex = index - 1
+                                if (targetIndex >= 0) {
+                                    val target = state.messages[targetIndex]
+                                    val targetRequester = requesterFor(target)
+                                    scope.launch {
+                                        listState.scrollToItem(targetIndex)
+                                        snapshotFlow {
+                                            listState.layoutInfo.visibleItemsInfo
+                                                .any { it.index == targetIndex }
+                                        }.first { it }
+                                        targetRequester.requestFocus()
+                                    }
+                                } else {
+                                    textFieldFocusRequester.requestFocus()
+                                }
+                            }
                         }
                     }
                 )
@@ -223,6 +258,7 @@ private fun ChatContent(
 
 @Composable
 private fun MessageBubble(
+    itemFocusRequester: FocusRequester,
     message: Message,
     onParseMessage: (String) -> List<TextToken>,
     onClickLink: (String) -> Unit,
@@ -237,7 +273,6 @@ private fun MessageBubble(
     val annotatedString = remember(tokens) {
         buildText(isTv, tokens, primaryColor, focusedColor, onClickLink)
     }
-    val itemFocusRequester = remember { FocusRequester() }
     var dialog by remember { mutableStateOf(false) }
     var boxModifier = Modifier.widthIn(max = 280.dp)
     if (isTv) {
@@ -359,17 +394,14 @@ private fun ScrollableMessage(
                             coroutineScope.launch { scrollState.animateScrollBy(scrollStep) }
                             true
                         }
-
                         Key.DirectionUp -> {
                             coroutineScope.launch { scrollState.animateScrollBy(-scrollStep) }
                             true
                         }
-
                         Key.Enter, Key.DirectionCenter -> {
                             onEnter()
                             true
                         }
-
                         else -> false
                     }
                 } else {