scext-onedrive-transfer/src/main/kotlin/Services.kt

package net.lamgc.scext.onedrive_transfer

import com.microsoft.aad.msal4j.ClientCredentialFactory
import com.microsoft.aad.msal4j.ConfidentialClientApplication
import com.microsoft.aad.msal4j.IAccount
import com.microsoft.aad.msal4j.SilentParameters
import com.microsoft.graph.authentication.IAuthenticationProvider
import com.microsoft.graph.httpcore.HttpClients
import com.microsoft.graph.models.Drive
import com.microsoft.graph.requests.GraphServiceClient
import mu.KotlinLogging
import okhttp3.Request
import org.jetbrains.exposed.sql.Database
import org.jetbrains.exposed.sql.SchemaUtils
import org.jetbrains.exposed.sql.transactions.transaction
import org.telegram.abilitybots.api.bot.BaseAbilityBot
import org.telegram.telegrambots.meta.api.objects.Document
import java.net.URL
import java.util.concurrent.CompletableFuture


class OneDriveTransferService(
    private val bot: BaseAbilityBot,
    private val config: ExtensionConfig,
    private val db: Database
) {
    private val logger = KotlinLogging.logger {  }

    val accountManager: OneDriveTransferSettingManager
    private val authClient: ConfidentialClientApplication = ConfidentialClientApplication.builder(
        config.clientId,
        ClientCredentialFactory.createFromSecret(config.clientSecret),
    )
        .authority(OneDriveTransferSettingManager.AUTHORITY)
        .setTokenCacheAccessAspect(DatabaseTokenCache(db))
        .build()

    init {
        transaction(db) {
            SchemaUtils.create(OneDriveTransferSettings, TokenCaches)
        }
        accountManager = OneDriveTransferSettingManager(authClient, db)
    }

    fun createGraphClient(userId: Long): GraphServiceClient<Request> {
        val cache = THREAD_CURRENT_GRAPH_CLIENT.get()
        if (cache?.tgUserId == userId) {
            return cache.client
        }
        val serviceClient = accountManager.getTransferSetting(userId)?.let {
            authClient.accounts.get().firstOrNull { account ->
                account.homeAccountId() == it.accountId
            }
        }?.let { Companion.createGraphClient(authClient, it) } ?: throw OneDriveNotLoginException()
        THREAD_CURRENT_GRAPH_CLIENT.set(ClientCache(userId, serviceClient))
        return serviceClient
    }

    fun createLoginUrl(userId: Long) =
        accountManager.createAuthorizationRequest(userId)

    fun updateAccount(userId: Long, redirectUrl: URL) =
        accountManager.updateAccount(userId, redirectUrl)

    fun listDriversByUserId(userId: Long): List<Drive> {
        val sites = try {
            val sites = createGraphClient(userId).sites()
                .buildRequest().get()?.currentPage ?: emptyList()
            sites.map {
                it.drive
            }.filterNotNull().toList()
        } catch (e: Exception) {
            logger.debug(e) { "获取 OneDrive 站点失败, 可能是用户没有权限或不是组织账号." }
            emptyList()
        }
        val drive = createGraphClient(userId).me().drive().buildRequest().get()
        val drives = createGraphClient(userId).drives()
            .buildRequest()
            .get()?.currentPage ?: emptyList()
        return mutableListOf<Drive>().apply {
            if (drive != null) {
                add(drive)
            }
            addAll(drives)
            addAll(sites)
        }
    }

    fun setDrive(userId: Long, driveId: String) {
        val transferSetting =
            accountManager.getTransferSetting(userId) ?: throw OneDriveNotLoginException()
        accountManager.doSomething {
            transferSetting.driveId = driveId
        }
    }

    fun getCurrentDrive(userId: Long): Drive? {
        val transferSetting =
            accountManager.getTransferSetting(userId) ?: throw OneDriveNotLoginException()
        val graphClient = createGraphClient(userId)
        return graphClient.drives(transferSetting.driveId)
            .buildRequest()
            .get()
    }

    fun submitUploadDocumentTask(userId: Long, messageId: Int, document: Document) {
        val transferSetting =
            accountManager.getTransferSetting(userId) ?: throw OneDriveNotLoginException()
        OneDriveTransferCenter.submitUploadTask(
            OneDriveTransferTask(
                userId,
                bot,
                this,
                document,
                transferSetting.driveId,
                transferSetting.storagePath,
            ).apply {
                extra["chatId"] = userId
                extra["messageId"] = messageId
            }
        )
    }

    companion object {
        private val THREAD_CURRENT_GRAPH_CLIENT = ThreadLocal<ClientCache>()

        fun createGraphClient(authClient: ConfidentialClientApplication, iAccount: IAccount): GraphServiceClient<Request> {
            return GraphServiceClient.builder()
                .httpClient(HttpClients.createDefault(MsalAuthorizationProvider(authClient, iAccount)))
                .buildClient()
        }
    }
}

private data class ClientCache(
    val tgUserId: Long,
    val client: GraphServiceClient<Request>,
)

class MsalAuthorizationProvider(private val authClientApplication: ConfidentialClientApplication, private val iAccount: IAccount) : IAuthenticationProvider {
    override fun getAuthorizationTokenAsync(requestUrl: URL): CompletableFuture<String> {
        return authClientApplication.acquireTokenSilently(
            SilentParameters.builder(OneDriveTransferSettingManager.OAUTH2_SCOPE, iAccount).build()
        ).thenApply { it.accessToken() }
    }

}