GoogleOAuthDesktop.kt

package com.codingwitharul.bill_it.common.auth

import com.codingwitharul.bill_it.common.model.GoogleUser
import com.codingwitharul.bill_it.data.networking.ApiClientHelper
import com.codingwitharul.bill_it.utils.toThrowable
import io.ktor.http.*
import io.ktor.server.cio.*
import io.ktor.server.engine.*
import io.ktor.client.request.*
import io.ktor.client.statement.*
import io.ktor.client.call.*
import io.ktor.server.response.respondText
import io.ktor.server.routing.get
import io.ktor.server.routing.routing
import kotlinx.coroutines.*
import java.awt.Desktop
import java.net.URI
import java.util.*
import java.util.concurrent.TimeUnit
import kotlin.time.Duration
import kotlin.time.Duration.Companion.minutes

class GoogleOAuthDesktop(
    private val clientId: String,
    private val clientSecret: String,
    private val apiClientHelper: ApiClientHelper,
    private val scopes: List<String> = listOf(
        "https://www.googleapis.com/auth/userinfo.profile",
        "https://www.googleapis.com/auth/userinfo.email",
        "https://www.googleapis.com/auth/user.phonenumbers.read"
    ),
    private val redirectHost: String = "localhost",
    private val redirectPort: Int = 9004,
    private val timeout: Duration = 1.minutes,
) {
    private val redirectUri = "http://$redirectHost:$redirectPort/callback"

    private val tokenUrl = "https://oauth2.googleapis.com/token"
    private val userInfoUrl = "https://www.googleapis.com/oauth2/v2/userinfo"


    suspend fun authenticate(): Result<GoogleUser?> {

        val state = UUID.randomUUID().toString()
        val authUrl = buildAuthUrl(state)
        Desktop.getDesktop().browse(URI(authUrl))
        val response = waitForAuthorizationCode(state)
        return response.fold(
            onSuccess = { tokenResponse ->
                if (tokenResponse != null) {
                    val googleUserProfile = fetchUserInfo(tokenResponse.accessToken)
                    if (googleUserProfile != null) {
                        Result.success(GoogleUser(
                            tokenResponse.accessToken,
                            googleUserProfile.id,
                            googleUserProfile.name,
                            googleUserProfile.pictureUrl,
                            googleUserProfile.familyName,
                            googleUserProfile.name,
                            tokenResponse.refreshToken
                        ))
                    } else {
                        Result.failure("Failed to fetch user profile.".toThrowable())
                    }
                } else {
                    Result.failure("Token response is null.".toThrowable())
                }
            },
            onFailure = {
                Result.failure(it)
            }
        )
    }

    private fun buildAuthUrl(state: String): String {
        val scopeEncoded = scopes.joinToString("+") { URI.create(it).toString() }
        return buildString {
            append("https://accounts.google.com/o/oauth2/v2/auth?")
            append("response_type=code&")
            append("client_id=$clientId&")
            append("redirect_uri=$redirectUri&")
            append("scope=$scopeEncoded&")
            append("state=$state&")
            append("access_type=offline&")
            append("prompt=consent")
        }
    }

    @OptIn(DelicateCoroutinesApi::class)
    private suspend fun waitForAuthorizationCode(expectedState: String): Result<GoogleTokenResponse?> {
        var authCode: String? = null
        val server = embeddedServer(CIO, port = redirectPort) {
            routing {
                get("/callback") {
                    val query = call.request.queryParameters
                    val code = query["code"]
                    val state = query["state"]

                    if (code != null && state == expectedState) {
                        authCode = code
                        call.respondText("Login successful. You may close this window.")
                        GlobalScope.launch {
                            delay(2000)
                            this@embeddedServer.engine.stop(gracePeriod = 1000L, timeout = 1000L, timeUnit = TimeUnit.MILLISECONDS)
                        }
                    } else {
                        call.respondText("Login failed or state mismatch.")
                    }
                }
            }
        }
        server.start(wait = false)
        var timer = 0
        var timeoutExpiry = false
        while (authCode == null && timer < timeout.inWholeMilliseconds) {
            delay(500)
            timer += 500
            timeoutExpiry = true
        }
        if (timeoutExpiry && authCode == null) {
            return Result.failure("Timeout waiting for auth code.".toThrowable())
        }
        return Result.success(exchangeCodeForToken(authCode!!))
    }

    private suspend fun exchangeCodeForToken(code: String): GoogleTokenResponse? {

        val response: HttpResponse = apiClientHelper.client.post(tokenUrl) {
            setBody(
                Parameters.build {
                    append("code", code)
                    append("client_id", clientId)
                    append("client_secret", clientSecret)
                    append("redirect_uri", redirectUri)
                    append("grant_type", "authorization_code")
                }.formUrlEncode()
            )
            headers {
                append(HttpHeaders.ContentType, ContentType.Application.FormUrlEncoded.toString())
            }
        }
        return response.body()
    }

    private suspend fun fetchUserInfo(accessToken: String): GoogleUserProfile? {
        return apiClientHelper.client.get(userInfoUrl) {
            headers {
                append(HttpHeaders.Authorization, "Bearer $accessToken")
            }
        }.body()
    }
}