// RFC - 6120: chapter 6
// XEP-0388: Extensible SASL Profile
import Foundation

enum AuthorizationStep: Codable {
    case notAuthorized
    case inProgress
    case authorized
}

enum AuthorizationError: Error {
    case noSupportedMechanisms
    case channelBindError
    case mechanismError(String)
}

final class AuthorizationModule: XmppModule {
    let id = "Authorization module"

    private weak var storage: (any XMPPStorage)?
    private var mechanism: AuthorizationMechanism?

    init(_ storage: any XMPPStorage) {
        self.storage = storage
    }

    func reduce(oldState: ClientState, with event: Event) -> ClientState {
        var newState = oldState
        switch event {
        case .startAuth:
            newState.authorizationStep = .inProgress

        case .gotAuthError:
            newState.authorizationStep = .notAuthorized

        case .authDone:
            newState.authorizationStep = .authorized

        default:
            break
        }
        return newState
    }

    func process(state: ClientState, with event: Event) async -> Event? {
        switch event {
        case .xmlInbound(let xml):
            guard state.isSocketSecured else { return nil }
            switch (xml.name, state.authorizationStep) {
            case ("stream:features", .notAuthorized):
                let credentials = await storage?.getCredentialsByUUID(state.credentialsId)
                return await selectBestAuthMechanism(xml, state.allowPlainAuth, state, credentials)

            case (_, .inProgress):
                return .challengeAuth(xml)

            default:
                return nil
            }

        case .startAuth(let xml):
            return .xmlOutbound(xml)

        case .challengeAuth(let xml):
            return await mechanism?.challenge(xml: xml)

        case .authDone:
            mechanism = nil
            return nil

        default:
            return nil
        }
    }
}

private extension AuthorizationModule {
    var supportedChannelBindings: [String] {
        ["tls-exporter", "tls-server-end-point"]
    }

    func selectBestAuthMechanism(_ xml: XMLElement, _ isPlainAllowed: Bool, _ state: ClientState, _ creds: Credentials?) async -> Event {
        await withCheckedContinuation { continuation in
            var sasl1: [AuthorizationMechanismType] = []
            var channelBindings: [String] = []
            var sasl2: [AuthorizationMechanismType] = []
            var inlines: XMLElement?

            // parse features
            for element in xml.nodes {
                // extract sasl1 mechanisms
                if element.name == "mechanisms" && element.xmlns == "urn:ietf:params:xml:ns:xmpp-sasl" {
                    sasl1 = element.nodes
                        .compactMap { AuthorizationMechanismType(rawValue: $0.content ?? "") }
                        .sorted { $0.priority < $1.priority }
                }
                // extract channel bindings
                if element.name == "sasl-channel-binding" && element.xmlns == "urn:xmpp:sasl-cb:0" {
                    channelBindings = element.nodes
                        .compactMap { $0.attributes["type"] }
                        .filter { supportedChannelBindings.contains($0) }
                }
                // extract sasl2
                if element.name == "authentication" && element.xmlns == "urn:xmpp:sasl:2" {
                    // sasl2 mechanisms
                    sasl2 = element.nodes
                        .filter { $0.name == "mechanism" && $0.xmlns == "urn:xmpp:sasl:2" }
                        .compactMap { AuthorizationMechanismType(rawValue: $0.content ?? "") }
                        .sorted { $0.priority < $1.priority }

                    // sasl2 inlines
                    inlines = element.nodes.first(where: { $0.name == "inline" && $0.xmlns == "urn:xmpp:sasl:2" })
                }
            }

            // filter out PLAIN if needed
            if !isPlainAllowed {
                sasl1 = sasl1.filter { $0 != .plain }
                sasl2 = sasl2.filter { $0 != .plain }
            }
            if sasl1.isEmpty && sasl2.isEmpty {
                continuation.resume(returning: .gotAuthError(.noSupportedMechanisms))
            }

            // select best authorization way
            var best = sasl1.map { ($0, SaslType.sasl1) }
            for mechanism in sasl2 {
                if best.isEmpty {
                    best.insert((mechanism, SaslType.sasl2), at: 0)
                } else if mechanism.priority <= best[0].0.priority {
                    best.insert((mechanism, SaslType.sasl2), at: 0)
                }
            }
            let selected = best[0]

            // init mechanism and start auth
            mechanism = AuthorizationMechanismImpl(
                type: selected.0,
                jid: state.jid,
                credentials: creds,
                saslType: selected.1,
                userAgent: state.userAgent,
                channelBind: nil, // TODO: check channel binding and implement *-PLUS mechanisms
                inlines: nil // TODO: Implement inlines
            )
            let request = mechanism?.initRequest
            if let request {
                continuation.resume(returning: .startAuth(request))
            } else {
                continuation.resume(returning: .gotAuthError(.mechanismError("Init request is empty")))
            }
        }
    }
}