Please fix this code and make it fully functional: package com.hx.chathx

import android.os.Bundle
import androidx.activity.ComponentActivity
import androidx.activity.compose.setContent
import androidx.compose.foundation.layout.*
import androidx.compose.foundation.lazy.LazyColumn
import androidx.compose.foundation.lazy.items
import androidx.compose.material3.*
import androidx.compose.runtime.*
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.unit.dp
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import dagger.hilt.android.AndroidEntryPoint
import dagger.hilt.android.lifecycle.HiltViewModel
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.launch
import okhttp3.*
import com.google.gson.Gson
import com.google.gson.JsonElement
import javax.inject.Inject
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext

@AndroidEntryPoint
class MainActivity : ComponentActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContent {
            ChatHXTheme {
                ChatScreen()
            }
        }
    }
}

@HiltViewModel
class ChatViewModel @Inject constructor(
    private val chatRepository: ChatRepository
) : ViewModel() {
    private val _uiState = MutableStateFlow(ChatUiState())
    val uiState: StateFlow<ChatUiState> = _uiState

    fun sendMessage(message: String) {
        viewModelScope.launch {
            _uiState.value = _uiState.value.copy(isLoading = true)
            try {
                val response = chatRepository.sendMessage(message)
                _uiState.value = _uiState.value.copy(
                    messages = _uiState.value.messages + ChatMessage(message, true) +
                            ChatMessage(response.answer, false) +
                            response.followUpQuestions.map { ChatMessage(it, false) },
                    isLoading = false
                )
            } catch (e: Exception) {
                _uiState.value = _uiState.value.copy(
                    error = e.message ?: "An error occurred",
                    isLoading = false
                )
            }
        }
    }
}

data class ChatUiState(
    val messages: List<ChatMessage> = emptyList(),
    val isLoading: Boolean = false,
    val error: String? = null
)

data class ChatMessage(val text: String, val isUser: Boolean)

data class ChatResponse(val answer: String, val followUpQuestions: List<String>)

class ChatRepository @Inject constructor(
    private val client: OkHttpClient,
    private val gson: Gson
) {
    suspend fun sendMessage(message: String): ChatResponse = withContext(Dispatchers.IO) {
        val url = "http://10.0.2.2:5000/api/chat"
        val requestBody = RequestBody.create(
            MediaType.parse("application/json"),
            gson.toJson(mapOf("message" to message))
        )
        val request = Request.Builder()
            .url(url)
            .post(requestBody)
            .build()

        val response = client.newCall(request).execute()
        if (!response.isSuccessful) throw Exception("Server error: ${response.code}")

        val responseBody = response.body()?.string() ?: throw Exception("Empty response body")
        val responseJson: JsonElement = gson.fromJson(responseBody, JsonElement::class.java)
        val answer = responseJson.asJsonObject.get("answer").asString
        val questions = responseJson.asJsonObject.get("follow_up_questions").asJsonArray.map { it.asString }

        ChatResponse(answer, questions)
    }
}

@Composable
fun ChatScreen(viewModel: ChatViewModel = hiltViewModel()) {
    val uiState by viewModel.uiState.collectAsState()

    Column(
        modifier = Modifier.fillMaxSize(),
        verticalArrangement = Arrangement.Bottom
    ) {
        ChatMessageArea(uiState.messages)
        UserInputArea(
            onSendMessage = viewModel::sendMessage,
            isLoading = uiState.isLoading
        )
    }

    uiState.error?.let { ErrorDialog(it) }
}

@Composable
fun ChatMessageArea(messages: List<ChatMessage>) {
    LazyColumn(
        modifier = Modifier
            .fillMaxWidth()
            .weight(1f),
        contentPadding = PaddingValues(16.dp)
    ) {
        items(messages) { message ->
            MessageCard(message)
        }
    }
}

@Composable
fun UserInputArea(
    onSendMessage: (String) -> Unit,
    isLoading: Boolean
) {
    var userInput by remember { mutableStateOf("") }

    Row(
        modifier = Modifier
            .fillMaxWidth()
            .padding(16.dp),
        verticalAlignment = Alignment.CenterVertically
    ) {
        OutlinedTextField(
            value = userInput,
            onValueChange = { userInput = it },
            label = { Text("Type your question...") },
            modifier = Modifier
                .weight(1f)
                .padding(end = 8.dp)
        )

        Button(
            onClick = {
                onSendMessage(userInput)
                userInput = ""
            },
            modifier = Modifier.wrapContentSize(),
            enabled = !isLoading
        ) {
            if (isLoading) {
                CircularProgressIndicator(color = Color.White)
            } else {
                Text("Send")
            }
        }
    }
}

@Composable
fun MessageCard(message: ChatMessage) {
    val alignment = if (message.isUser) Alignment.End else Alignment.Start
    val backgroundColor = if (message.isUser) Color.Blue else Color.LightGray
    val textColor = if (message.isUser) Color.White else Color.Black

    Row(
        modifier = Modifier
            .fillMaxWidth()
            .padding(8.dp),
        horizontalArrangement = if (message.isUser) Arrangement.End else Arrangement.Start
    ) {
        Text(
            text = message.text,
            color = textColor,
            fontWeight = FontWeight.Bold,
            modifier = Modifier
                .padding(8.dp)
                .background(backgroundColor)
        )
    }
}

@Composable
fun ErrorDialog(errorMessage: String) {
    AlertDialog(
        onDismissRequest = { /* Dismiss logic */ },
        title = { Text("Error") },
        text = { Text(errorMessage) },
        confirmButton = {
            Button(onClick = { /* Dismiss logic */ }) {
                Text("OK")
            }
        }
    )
}

@Composable
fun ChatHXTheme(content: @Composable () -> Unit) {
    MaterialTheme(
        colorScheme = darkColorScheme(
            primary = Color(0xFF1976D2),
            surface = Color(0xFF121212)
        )
    ) {
        Surface(
            modifier = Modifier.fillMaxSize(),
            color = MaterialTheme.colorScheme.background
        ) {
            content()
        }
    }
}