package dev.zxq5.chatapp.android.core.data

import android.content.Context
import android.content.SharedPreferences
import android.util.Base64
import android.util.Log
import androidx.core.content.edit
import androidx.security.crypto.EncryptedSharedPreferences
import androidx.security.crypto.MasterKey
import org.json.JSONObject
import java.time.Instant

private const val KEY = "auth_token"
private const val TWOFA_KEY = "twofa_enabled"

// In your ChatClient.kt or a dedicated TokenStore
class TokenStore(appContext: Context) {
    private val context = appContext.applicationContext;

    private fun prefs(): SharedPreferences {
        return EncryptedSharedPreferences.create(
            context,
            "secure_prefs",
            MasterKey.Builder(context)
                .setKeyScheme(MasterKey.KeyScheme.AES256_GCM)
                .build(),
            EncryptedSharedPreferences.PrefKeyEncryptionScheme.AES256_SIV,
            EncryptedSharedPreferences.PrefValueEncryptionScheme.AES256_GCM
        )
    }

    fun save(token: String) {
        Log.d("TokenStore", "Saving token: $token")

        prefs().edit { putString(KEY, token) }
    }

    fun get(): String? {
        val ret = prefs().getString(KEY, null)
        Log.d("TokenStore", "Retrieved token: $ret")
        return ret
    }

    fun isExpired(): Boolean {
        val token = get() ?: return true
        return try {
            val payload = token.split(".")[1]
            val padded = payload + "==".take((4 - payload.length % 4) % 4)
            val jsonString = String(Base64.decode(padded, Base64.URL_SAFE))
            val json = JSONObject(jsonString)
            if (json.has("exp")) {
                val exp = json.getLong("exp")
                val now = Instant.now().epochSecond
                now >= exp
            } else {
                false // If no exp claim, assume not expired or handle differently
            }
        } catch (e: Exception) {
            true // If we can't parse it, treat it as expired
        }
    }


    fun save2faEnabled( enabled: Boolean) =
        prefs().edit { putBoolean(TWOFA_KEY, enabled) }

    fun is2faEnabled(): Boolean =
        prefs().getBoolean(TWOFA_KEY, false)

    fun clear() =
        prefs().edit { remove(KEY).remove(TWOFA_KEY) }

    fun getUserId(): Int? {
        val token = get() ?: return null
        return getUserIdFromToken(token)
    }
}

fun getUserIdFromToken(token: String): Int? {
    return try {
        val payload = token.split(".")[1]
        // base64url needs padding restored
        val padded = payload + "==".take((4 - payload.length % 4) % 4)
        val jsonString = String(Base64.decode(padded, Base64.URL_SAFE))
        val json = JSONObject(jsonString)

        // Handle both standard 'sub' and custom 'user_id'
        when {
            json.has("sub") -> json.getInt("sub")
            json.has("user_id") -> json.getInt("user_id")
            else -> null
        }
    } catch (e: Exception) {
        null
    }
}

fun getScopeFromToken(token: String): String? {
    return try {
        val payload = token.split(".")[1]
        val padded = payload + "==".take((4 - payload.length % 4) % 4)
        val jsonString = String(Base64.decode(padded, Base64.URL_SAFE))
        val json = JSONObject(jsonString)
        if (json.has("scope")) json.getString("scope") else null
    } catch (e: Exception) {
        null
    }
}