mobile-android/app/src/main/java/com/rosetta/messenger/providers/AuthState.kt

package com.rosetta.messenger.providers

import android.content.Context
import androidx.compose.runtime.*
import com.rosetta.messenger.data.AccountManager
import com.rosetta.messenger.data.DecryptedAccount
import com.rosetta.messenger.data.EncryptedAccount
import com.rosetta.messenger.crypto.CryptoManager
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.*
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext

/**
 * Auth state management - similar to AuthContext in React Native version
 */
sealed class AuthStatus {
    object Loading : AuthStatus()
    object Unauthenticated : AuthStatus()
    data class Authenticated(val account: DecryptedAccount) : AuthStatus()
    data class Locked(val publicKey: String) : AuthStatus()
}

data class AuthStateData(
    val status: AuthStatus = AuthStatus.Loading,
    val accounts: List<EncryptedAccount> = emptyList(),
    val hasExistingAccounts: Boolean = false
)

class AuthStateManager(
    private val context: Context,
    private val scope: CoroutineScope
) {
    private val accountManager = AccountManager(context)
    
    private val _state = MutableStateFlow(AuthStateData())
    val state: StateFlow<AuthStateData> = _state.asStateFlow()
    
    init {
        scope.launch {
            loadAccounts()
            checkAuthStatus()
        }
    }
    
    private suspend fun loadAccounts() = withContext(Dispatchers.IO) {
        val accounts = accountManager.getAllAccounts()
        _state.update { it.copy(
            accounts = accounts,
            hasExistingAccounts = accounts.isNotEmpty()
        )}
    }
    
    private suspend fun checkAuthStatus() {
        accountManager.isLoggedIn.collect { isLoggedIn ->
            if (isLoggedIn) {
                accountManager.currentPublicKey.first()?.let { publicKey ->
                    _state.update { it.copy(
                        status = AuthStatus.Locked(publicKey)
                    )}
                }
            } else {
                _state.update { it.copy(
                    status = AuthStatus.Unauthenticated
                )}
            }
        }
    }
    
    /**
     * Create new account from seed phrase
     */
    suspend fun createAccount(
        seedPhrase: List<String>,
        password: String,
        name: String = "Rosetta Account"
    ): Result<DecryptedAccount> = withContext(Dispatchers.Default) {
        try {
            val keyPair = CryptoManager.generateKeyPairFromSeed(seedPhrase)
            val privateKeyHash = CryptoManager.generatePrivateKeyHash(keyPair.privateKey)
            
            val encryptedPrivateKey = CryptoManager.encryptWithPassword(
                keyPair.privateKey, password
            )
            val encryptedSeedPhrase = CryptoManager.encryptWithPassword(
                seedPhrase.joinToString(" "), password
            )
            
            withContext(Dispatchers.IO) {
                val encryptedAccount = EncryptedAccount(
                    publicKey = keyPair.publicKey,
                    encryptedPrivateKey = encryptedPrivateKey,
                    encryptedSeedPhrase = encryptedSeedPhrase,
                    name = name
                )
                accountManager.saveAccount(encryptedAccount)
                accountManager.setCurrentAccount(keyPair.publicKey)
            }
            
            val decryptedAccount = DecryptedAccount(
                publicKey = keyPair.publicKey,
                privateKey = keyPair.privateKey,
                seedPhrase = seedPhrase,
                privateKeyHash = privateKeyHash,
                name = name
            )
            
            _state.update { it.copy(
                status = AuthStatus.Authenticated(decryptedAccount)
            )}
            
            loadAccounts()
            Result.success(decryptedAccount)
        } catch (e: Exception) {
            Result.failure(e)
        }
    }
    
    /**
     * Import existing account from seed phrase
     */
    suspend fun importAccount(
        seedPhrase: List<String>,
        password: String,
        name: String = "Imported Account"
    ): Result<DecryptedAccount> {
        return createAccount(seedPhrase, password, name)
    }
    
    /**
     * Unlock account with password
     */
    suspend fun unlock(
        publicKey: String,
        password: String
    ): Result<DecryptedAccount> = withContext(Dispatchers.Default) {
        try {
            val encryptedAccount = withContext(Dispatchers.IO) {
                accountManager.getAccount(publicKey)
            } ?: return@withContext Result.failure(Exception("Account not found"))
            
            val privateKey = CryptoManager.decryptWithPassword(
                encryptedAccount.encryptedPrivateKey, password
            ) ?: return@withContext Result.failure(Exception("Invalid password"))
            
            val seedPhraseStr = CryptoManager.decryptWithPassword(
                encryptedAccount.encryptedSeedPhrase, password
            ) ?: return@withContext Result.failure(Exception("Invalid password"))
            
            val seedPhrase = seedPhraseStr.split(" ")
            val keyPair = CryptoManager.generateKeyPairFromSeed(seedPhrase)
            
            if (keyPair.publicKey != publicKey) {
                return@withContext Result.failure(Exception("Invalid password or corrupted data"))
            }
            
            val privateKeyHash = CryptoManager.generatePrivateKeyHash(privateKey)
            
            val decryptedAccount = DecryptedAccount(
                publicKey = publicKey,
                privateKey = privateKey,
                seedPhrase = seedPhrase,
                privateKeyHash = privateKeyHash,
                name = encryptedAccount.name
            )
            
            withContext(Dispatchers.IO) {
                accountManager.setCurrentAccount(publicKey)
            }
            
            _state.update { it.copy(
                status = AuthStatus.Authenticated(decryptedAccount)
            )}
            
            Result.success(decryptedAccount)
        } catch (e: Exception) {
            Result.failure(e)
        }
    }
    
    /**
     * Lock account
     */
    fun lock() {
        val currentStatus = _state.value.status
        if (currentStatus is AuthStatus.Authenticated) {
            _state.update { it.copy(
                status = AuthStatus.Locked(currentStatus.account.publicKey)
            )}
        }
    }
    
    /**
     * Logout
     */
    suspend fun logout() {
        withContext(Dispatchers.IO) {
            accountManager.logout()
        }
        _state.update { it.copy(
            status = AuthStatus.Unauthenticated
        )}
        loadAccounts()
    }
    
    /**
     * Switch to different account
     */
    fun switchAccount(publicKey: String) {
        _state.update { it.copy(
            status = AuthStatus.Locked(publicKey)
        )}
    }
    
    /**
     * Delete account permanently
     */
    suspend fun deleteAccount(publicKey: String) = withContext(Dispatchers.IO) {
        val accounts = accountManager.getAllAccounts().toMutableList()
        accounts.removeAll { it.publicKey == publicKey }
        
        accountManager.clearAll()
        accounts.forEach { accountManager.saveAccount(it) }
        
        loadAccounts()
        
        val currentStatus = _state.value.status
        if ((currentStatus is AuthStatus.Authenticated && currentStatus.account.publicKey == publicKey) ||
            (currentStatus is AuthStatus.Locked && currentStatus.publicKey == publicKey)) {
            logout()
        }
    }
}

@Composable
fun rememberAuthState(context: Context): AuthStateManager {
    val scope = rememberCoroutineScope()
    return remember(context) {
        AuthStateManager(context, scope)
    }
}

@Composable
fun ProvideAuthState(
    authState: AuthStateManager,
    content: @Composable (AuthStateData) -> Unit
) {
    val state by authState.state.collectAsState()
    content(state)
}