classroom_app/classroom_app/Auth/GoogleOAuthService.swift

import Foundation
import CryptoKit
import AppKit
import Network
import WebKit

struct GoogleOAuthTokens: Codable, Equatable {
    var accessToken: String
    var refreshToken: String?
    var expiresAt: Date
    var scope: String?
    var tokenType: String?
}

struct GoogleUserProfile: Codable, Equatable {
    var name: String?
    var email: String?
    var picture: String?
}

enum GoogleOAuthError: Error {
    case missingClientId
    case missingClientSecret
    case invalidCallbackURL
    case missingAuthorizationCode
    case tokenExchangeFailed(String)
    case unableToOpenBrowser
    case authenticationTimedOut
    case noStoredTokens
}

final class GoogleOAuthService: NSObject {
    static let shared = GoogleOAuthService()
    private var inAppOAuthWindowController: InAppOAuthWindowController?

    // Stored in UserDefaults so you can configure without rebuilding.
    // Put your OAuth Desktop client ID here (from Google Cloud Console).
    private let clientIdDefaultsKey = "google.oauth.clientId"
    private let clientSecretDefaultsKey = "google.oauth.clientSecret"
    private let bundledClientId = "1027920783127-tu96fn69edr6fimn32nqh9rfj711fm2i.apps.googleusercontent.com"
    private let bundledClientSecret = "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"
    private let legacyMeetingClientId = "824412072260-m9g5g6mlemnb0o079rtuqnh0e1unmelc.apps.googleusercontent.com"

    // Profile/email make login feel complete in-app.
    private let scopes = [
        "openid",
        "email",
        "profile",
        // Classroom To-do (assignments/quizzes)
        "https://www.googleapis.com/auth/classroom.courses.readonly",
        "https://www.googleapis.com/auth/classroom.courses",
        "https://www.googleapis.com/auth/classroom.coursework.me.readonly",
        "https://www.googleapis.com/auth/classroom.announcements.readonly",
        "https://www.googleapis.com/auth/classroom.rosters.readonly",
        "https://www.googleapis.com/auth/classroom.rosters"
    ]

    private let tokenStore = KeychainTokenStore()
    private override init() {
        super.init()
        migrateLegacyOAuthOverridesIfNeeded()
    }

    func configuredClientId() -> String? {
        let value = UserDefaults.standard.string(forKey: clientIdDefaultsKey)?.trimmingCharacters(in: .whitespacesAndNewlines)
        if let value, value.isEmpty == false { return value }
        return bundledClientId
    }

    func setClientIdForTesting(_ clientId: String) {
        UserDefaults.standard.set(clientId, forKey: clientIdDefaultsKey)
    }

    func configuredClientSecret() -> String? {
        let value = UserDefaults.standard.string(forKey: clientSecretDefaultsKey)?.trimmingCharacters(in: .whitespacesAndNewlines)
        if let value, value.isEmpty == false { return value }
        return bundledClientSecret
    }

    func setClientSecretForTesting(_ clientSecret: String) {
        UserDefaults.standard.set(clientSecret, forKey: clientSecretDefaultsKey)
    }

    private func migrateLegacyOAuthOverridesIfNeeded() {
        let configuredId = UserDefaults.standard.string(forKey: clientIdDefaultsKey)?.trimmingCharacters(in: .whitespacesAndNewlines)
        guard configuredId == legacyMeetingClientId else { return }
        // Remove old project override so this app uses the Classroom OAuth client.
        UserDefaults.standard.removeObject(forKey: clientIdDefaultsKey)
        UserDefaults.standard.removeObject(forKey: clientSecretDefaultsKey)
    }

    func signOut() throws {
        try tokenStore.deleteTokens()
    }

    func loadTokens() -> GoogleOAuthTokens? {
        try? tokenStore.readTokens()
    }

    func fetchUserProfile(accessToken: String) async throws -> GoogleUserProfile {
        var request = URLRequest(url: URL(string: "https://openidconnect.googleapis.com/v1/userinfo")!)
        request.httpMethod = "GET"
        request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")

        let (data, response) = try await URLSession.shared.data(for: request)
        guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            let details = String(data: data, encoding: .utf8) ?? "HTTP \((response as? HTTPURLResponse)?.statusCode ?? -1)"
            throw GoogleOAuthError.tokenExchangeFailed(details)
        }
        return try JSONDecoder().decode(GoogleUserProfile.self, from: data)
    }

    func validAccessToken(presentingWindow: NSWindow?) async throws -> String {
        if var tokens = try tokenStore.readTokens() {
            if tokens.expiresAt.timeIntervalSinceNow > 60 {
                return tokens.accessToken
            }
            if let refreshed = try await refreshTokens(tokens) {
                tokens = refreshed
                try tokenStore.writeTokens(tokens)
                return tokens.accessToken
            }
        }

        let tokens = try await interactiveSignIn(presentingWindow: presentingWindow)
        try tokenStore.writeTokens(tokens)
        return tokens.accessToken
    }

    // MARK: - Interactive sign-in (Authorization Code + PKCE)

    private func interactiveSignIn(presentingWindow: NSWindow?) async throws -> GoogleOAuthTokens {
        _ = presentingWindow
        guard let clientId = configuredClientId() else { throw GoogleOAuthError.missingClientId }
        guard let clientSecret = configuredClientSecret() else { throw GoogleOAuthError.missingClientSecret }
        let codeVerifier = Self.randomURLSafeString(length: 64)
        let codeChallenge = Self.pkceChallenge(for: codeVerifier)
        let state = Self.randomURLSafeString(length: 32)

        let loopback = try await OAuthLoopbackServer.start()
        defer { loopback.stop() }
        let redirectURI = loopback.redirectURI

        var components = URLComponents(string: "https://accounts.google.com/o/oauth2/v2/auth")!
        components.queryItems = [
            URLQueryItem(name: "client_id", value: clientId),
            URLQueryItem(name: "redirect_uri", value: redirectURI),
            URLQueryItem(name: "response_type", value: "code"),
            URLQueryItem(name: "scope", value: scopes.joined(separator: " ")),
            URLQueryItem(name: "state", value: state),
            URLQueryItem(name: "code_challenge", value: codeChallenge),
            URLQueryItem(name: "code_challenge_method", value: "S256"),
            URLQueryItem(name: "access_type", value: "offline"),
            URLQueryItem(name: "prompt", value: "consent")
        ]

        guard let authURL = components.url else { throw GoogleOAuthError.invalidCallbackURL }
        let opened = await MainActor.run { [self] in
            openAuthURLInApp(authURL)
        }
        guard opened else { throw GoogleOAuthError.unableToOpenBrowser }
        defer {
            Task { @MainActor [weak self] in
                self?.closeInAppOAuthWindow()
            }
        }
        let callbackURL = try await loopback.waitForCallback()

        guard let returnedState = URLComponents(url: callbackURL, resolvingAgainstBaseURL: false)?
            .queryItems?.first(where: { $0.name == "state" })?.value,
              returnedState == state else {
            throw GoogleOAuthError.invalidCallbackURL
        }

        guard let code = URLComponents(url: callbackURL, resolvingAgainstBaseURL: false)?
            .queryItems?.first(where: { $0.name == "code" })?.value,
              code.isEmpty == false else {
            throw GoogleOAuthError.missingAuthorizationCode
        }

        return try await exchangeCodeForTokens(
            code: code,
            codeVerifier: codeVerifier,
            redirectURI: redirectURI,
            clientId: clientId,
            clientSecret: clientSecret
        )
    }

    @MainActor
    private func openAuthURLInApp(_ url: URL) -> Bool {
        let controller: InAppOAuthWindowController
        if let existing = inAppOAuthWindowController {
            controller = existing
        } else {
            controller = InAppOAuthWindowController()
            inAppOAuthWindowController = controller
        }
        controller.load(url: url)
        controller.showWindow(nil)
        controller.window?.makeKeyAndOrderFront(nil)
        controller.window?.orderFrontRegardless()
        NSApp.activate(ignoringOtherApps: true)
        return true
    }

    @MainActor
    private func closeInAppOAuthWindow() {
        inAppOAuthWindowController?.close()
        inAppOAuthWindowController = nil
    }

    private func exchangeCodeForTokens(code: String, codeVerifier: String, redirectURI: String, clientId: String, clientSecret: String) async throws -> GoogleOAuthTokens {
        var request = URLRequest(url: URL(string: "https://oauth2.googleapis.com/token")!)
        request.httpMethod = "POST"
        request.setValue("application/x-www-form-urlencoded; charset=utf-8", forHTTPHeaderField: "Content-Type")
        request.httpBody = Self.formURLEncoded([
            "client_id": clientId,
            "client_secret": clientSecret,
            "code": code,
            "code_verifier": codeVerifier,
            "redirect_uri": redirectURI,
            "grant_type": "authorization_code"
        ])

        let (data, response) = try await URLSession.shared.data(for: request)
        guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            let details = String(data: data, encoding: .utf8) ?? "HTTP \((response as? HTTPURLResponse)?.statusCode ?? -1)"
            throw GoogleOAuthError.tokenExchangeFailed(details)
        }

        struct TokenResponse: Decodable {
            let access_token: String
            let expires_in: Double
            let refresh_token: String?
            let scope: String?
            let token_type: String?
        }

        let decoded = try JSONDecoder().decode(TokenResponse.self, from: data)
        return GoogleOAuthTokens(
            accessToken: decoded.access_token,
            refreshToken: decoded.refresh_token,
            expiresAt: Date().addingTimeInterval(decoded.expires_in),
            scope: decoded.scope,
            tokenType: decoded.token_type
        )
    }

    private func refreshTokens(_ tokens: GoogleOAuthTokens) async throws -> GoogleOAuthTokens? {
        guard let refreshToken = tokens.refreshToken else { return nil }
        guard let clientId = configuredClientId() else { throw GoogleOAuthError.missingClientId }
        guard let clientSecret = configuredClientSecret() else { throw GoogleOAuthError.missingClientSecret }

        var request = URLRequest(url: URL(string: "https://oauth2.googleapis.com/token")!)
        request.httpMethod = "POST"
        request.setValue("application/x-www-form-urlencoded; charset=utf-8", forHTTPHeaderField: "Content-Type")
        request.httpBody = Self.formURLEncoded([
            "client_id": clientId,
            "client_secret": clientSecret,
            "refresh_token": refreshToken,
            "grant_type": "refresh_token"
        ])

        let (data, response) = try await URLSession.shared.data(for: request)
        guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            return nil
        }

        struct RefreshResponse: Decodable {
            let access_token: String
            let expires_in: Double
            let scope: String?
            let token_type: String?
        }

        let decoded = try JSONDecoder().decode(RefreshResponse.self, from: data)
        return GoogleOAuthTokens(
            accessToken: decoded.access_token,
            refreshToken: refreshToken,
            expiresAt: Date().addingTimeInterval(decoded.expires_in),
            scope: decoded.scope ?? tokens.scope,
            tokenType: decoded.token_type ?? tokens.tokenType
        )
    }

    // MARK: - Helpers

    private static func pkceChallenge(for verifier: String) -> String {
        let data = Data(verifier.utf8)
        let digest = SHA256.hash(data: data)
        return Data(digest).base64URLEncodedString()
    }

    private static func randomURLSafeString(length: Int) -> String {
        var bytes = [UInt8](repeating: 0, count: length)
        _ = SecRandomCopyBytes(kSecRandomDefault, bytes.count, &bytes)
        return Data(bytes).base64URLEncodedString()
    }

    private static func formURLEncoded(_ params: [String: String]) -> Data {
        let allowed = CharacterSet(charactersIn: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-._~")
        let pairs = params.map { key, value -> String in
            let k = key.addingPercentEncoding(withAllowedCharacters: allowed) ?? key
            let v = value.addingPercentEncoding(withAllowedCharacters: allowed) ?? value
            return "\(k)=\(v)"
        }
        .joined(separator: "&")
        return Data(pairs.utf8)
    }
}

private extension Data {
    func base64URLEncodedString() -> String {
        let s = base64EncodedString()
        return s
            .replacingOccurrences(of: "+", with: "-")
            .replacingOccurrences(of: "/", with: "_")
            .replacingOccurrences(of: "=", with: "")
    }
}

private final class OAuthLoopbackServer {
    private let queue = DispatchQueue(label: "google.oauth.loopback.server")
    private let listener: NWListener
    private var readyContinuation: CheckedContinuation<Void, Error>?
    private var callbackContinuation: CheckedContinuation<URL, Error>?
    private var callbackURL: URL?

    private init(listener: NWListener) {
        self.listener = listener
    }

    static func start() async throws -> OAuthLoopbackServer {
        let listener = try NWListener(using: .tcp, on: .any)
        let server = OAuthLoopbackServer(listener: listener)
        try await server.startListening()
        return server
    }

    var redirectURI: String {
        let port = listener.port?.rawValue ?? 0
        return "http://127.0.0.1:\(port)/oauth2redirect"
    }

    func waitForCallback(timeoutSeconds: Double = 120) async throws -> URL {
        try await withThrowingTaskGroup(of: URL.self) { group in
            group.addTask { [weak self] in
                guard let self else { throw GoogleOAuthError.invalidCallbackURL }
                return try await self.awaitCallback()
            }
            group.addTask {
                try await Task.sleep(nanoseconds: UInt64(timeoutSeconds * 1_000_000_000))
                throw GoogleOAuthError.authenticationTimedOut
            }

            let url = try await group.next()!
            group.cancelAll()
            return url
        }
    }

    func stop() {
        queue.async {
            self.listener.cancel()
            if let callbackContinuation = self.callbackContinuation {
                self.callbackContinuation = nil
                callbackContinuation.resume(throwing: GoogleOAuthError.authenticationTimedOut)
            }
        }
    }

    private func startListening() async throws {
        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
            queue.async {
                self.readyContinuation = continuation
                self.listener.stateUpdateHandler = { [weak self] state in
                    guard let self else { return }
                    switch state {
                    case .ready:
                        if let readyContinuation = self.readyContinuation {
                            self.readyContinuation = nil
                            readyContinuation.resume()
                        }
                    case .failed(let error):
                        if let readyContinuation = self.readyContinuation {
                            self.readyContinuation = nil
                            readyContinuation.resume(throwing: error)
                        }
                    case .cancelled:
                        if let readyContinuation = self.readyContinuation {
                            self.readyContinuation = nil
                            readyContinuation.resume(throwing: GoogleOAuthError.invalidCallbackURL)
                        }
                    default:
                        break
                    }
                }
                self.listener.newConnectionHandler = { [weak self] connection in
                    self?.handle(connection: connection)
                }
                self.listener.start(queue: self.queue)
            }
        }
    }

    private func awaitCallback() async throws -> URL {
        if let callbackURL { return callbackURL }
        return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<URL, Error>) in
            queue.async {
                if let callbackURL = self.callbackURL {
                    continuation.resume(returning: callbackURL)
                    return
                }
                self.callbackContinuation = continuation
            }
        }
    }

    private func handle(connection: NWConnection) {
        connection.start(queue: queue)
        connection.receive(minimumIncompleteLength: 1, maximumLength: 8192) { [weak self] data, _, _, _ in
            guard let self else { return }
            let requestLine = data
                .flatMap { String(data: $0, encoding: .utf8) }?
                .split(separator: "\r\n", omittingEmptySubsequences: false)
                .first
                .map(String.init)

            var parsedURL: URL?
            if let requestLine {
                let parts = requestLine.split(separator: " ")
                if parts.count >= 2 {
                    let pathAndQuery = String(parts[1])
                    parsedURL = URL(string: "http://127.0.0.1\(pathAndQuery)")
                }
            }

            self.sendHTTPResponse(connection: connection, success: parsedURL != nil)

            if let parsedURL {
                self.callbackURL = parsedURL
                DispatchQueue.main.async {
                    // Bring the app back to foreground once OAuth redirects.
                    NSApp.activate(ignoringOtherApps: true)
                }
                if let continuation = self.callbackContinuation {
                    self.callbackContinuation = nil
                    continuation.resume(returning: parsedURL)
                }
                self.listener.cancel()
            }
            connection.cancel()
        }
    }

    private func sendHTTPResponse(connection: NWConnection, success: Bool) {
        let body = success
            ? "<html><body><h3>Authentication complete</h3><p>You can return to the app.</p></body></html>"
            : "<html><body><h3>Authentication failed</h3></body></html>"
        let response = """
        HTTP/1.1 200 OK\r
        Content-Type: text/html; charset=utf-8\r
        Content-Length: \(body.utf8.count)\r
        Connection: close\r
        \r
        \(body)
        """
        connection.send(content: Data(response.utf8), completion: .contentProcessed { _ in })
    }
}

extension GoogleOAuthError: LocalizedError {
    var errorDescription: String? {
        switch self {
        case .missingClientId:
            return "Missing Google OAuth Client ID."
        case .missingClientSecret:
            return "Missing Google OAuth Client Secret."
        case .invalidCallbackURL:
            return "Invalid OAuth callback URL."
        case .missingAuthorizationCode:
            return "Google did not return an authorization code."
        case .tokenExchangeFailed(let details):
            return "Token exchange failed: \(details)"
        case .unableToOpenBrowser:
            return "Could not open browser for Google sign-in."
        case .authenticationTimedOut:
            return "Google sign-in timed out."
        case .noStoredTokens:
            return "No stored Google tokens found."
        }
    }
}

@MainActor
private final class OAuthWebViewContainerView: NSView {
    private let webView: WKWebView

    init(webView: WKWebView) {
        self.webView = webView
        super.init(frame: .zero)
        autoresizingMask = [.width, .height]
        addSubview(webView)
    }

    @available(*, unavailable)
    required init?(coder: NSCoder) {
        nil
    }

    override func layout() {
        super.layout()
        webView.frame = bounds
    }
}

@MainActor
private final class InAppOAuthWindowController: NSWindowController, WKNavigationDelegate, WKUIDelegate {
    private let webView: WKWebView

    init() {
        let config = WKWebViewConfiguration()
        if #available(macOS 11.0, *) {
            config.defaultWebpagePreferences.allowsContentJavaScript = true
        }
        self.webView = WKWebView(frame: .zero, configuration: config)
        // Frame-based layout avoids Auto Layout conflicts with WKWebView’s internal constraints.
        let container = OAuthWebViewContainerView(webView: webView)

        let window = NSWindow(
            contentRect: NSRect(x: 0, y: 0, width: 980, height: 760),
            styleMask: [.titled, .closable, .miniaturizable, .resizable],
            backing: .buffered,
            defer: false
        )
        window.title = "Google Sign-In"
        window.contentView = container
        window.center()
        super.init(window: window)
        webView.navigationDelegate = self
        webView.uiDelegate = self
    }

    @available(*, unavailable)
    required init?(coder: NSCoder) {
        nil
    }

    func load(url: URL) {
        webView.load(URLRequest(url: url))
    }

    private func shouldOpenURLExternally(_ url: URL) -> Bool {
        let scheme = (url.scheme ?? "").lowercased()
        guard !scheme.isEmpty else { return false }
        // Google auth popups can navigate through about:blank/javascript:
        // as an intermediate step. Don't hand these to NSWorkspace.
        return scheme != "about" && scheme != "javascript"
    }

    func webView(
        _ webView: WKWebView,
        decidePolicyFor navigationAction: WKNavigationAction,
        decisionHandler: @escaping (WKNavigationActionPolicy) -> Void
    ) {
        guard let url = navigationAction.request.url else {
            decisionHandler(.allow)
            return
        }
        let scheme = (url.scheme ?? "").lowercased()
        if scheme == "http" || scheme == "https" {
            decisionHandler(.allow)
            return
        }
        if shouldOpenURLExternally(url) {
            NSWorkspace.shared.open(url)
        }
        decisionHandler(.cancel)
    }

    func webView(
        _ webView: WKWebView,
        createWebViewWith configuration: WKWebViewConfiguration,
        for navigationAction: WKNavigationAction,
        windowFeatures: WKWindowFeatures
    ) -> WKWebView? {
        if navigationAction.targetFrame == nil, let requestURL = navigationAction.request.url {
            let scheme = (requestURL.scheme ?? "").lowercased()
            if scheme == "http" || scheme == "https" {
                webView.load(URLRequest(url: requestURL))
            } else {
                if shouldOpenURLExternally(requestURL) {
                    NSWorkspace.shared.open(requestURL)
                }
            }
        }
        return nil
    }
}