import Combine
import CommonCrypto
import CryptoKit
import Foundation
import Network

final class AuthMiddleware {
    fileprivate static let ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
    var clientNonce = ""

    func middleware(state: ClientState, action: ClientAction) -> AnyPublisher<ClientAction, Never> {
        switch action {
        case .authAction(.startAuth):
            // select most appropriate auth mechanism based on priority
            guard
                let mechanism = state.streamState.streamFeatures
                    .filter({ $0.featureName == "mechanism" })
                    .sorted(by: { $0.priority < $1.priority })
                    .first?.additionalInfo
            else {
                return Empty().eraseToAnyPublisher()
            }

            // start auth
            switch mechanism {
            case "PLAIN":
                let str = makePlainChallenge(jid: state.jid, pass: state.authState.pass ?? "")
                let req = XMLElement(
                    name: "auth",
                    xmlns: "urn:ietf:params:xml:ns:xmpp-sasl",
                    attributes: [
                        "mechanism": "PLAIN"
                    ],
                    content: str,
                    nodes: []
                )
                return Just(ClientAction.streamAction(.streamSend(req)))
                    .eraseToAnyPublisher()

            case "SCRAM-SHA-1":
                let (requestStr, clientNonce) = scramSha1InitRequestString(jid: state.jid)
                self.clientNonce = clientNonce
                let req = XMLElement(
                    name: "auth",
                    xmlns: "urn:ietf:params:xml:ns:xmpp-sasl",
                    attributes: ["mechanism": "SCRAM-SHA-1"],
                    content: requestStr,
                    nodes: []
                )
                return Just(ClientAction.streamAction(.streamSend(req)))
                    .eraseToAnyPublisher()

            default:
                return Empty().eraseToAnyPublisher()
            }

        case .parserAction(.elementParsed(let element)):
            if element.xmlns == "urn:ietf:params:xml:ns:xmpp-sasl" {
                switch element.name {
                case "success":
                    return Just(ClientAction.authAction(.gotAuthResult(.authorized)))
                        .eraseToAnyPublisher()

                case "challenge":
                    return Just(ClientAction.authAction(.gotAuthChallenge(element.content ?? "")))
                        .eraseToAnyPublisher()

                default:
                    return Just(ClientAction.authAction(.gotAuthResult(.failed)))
                        .eraseToAnyPublisher()
                }
            } else {
                return Empty().eraseToAnyPublisher()
            }

        case .authAction(.gotAuthChallenge(let msg)):
            guard
                let msgBytes = Foundation.Data(base64Encoded: msg, options: NSData.Base64DecodingOptions(rawValue: 0)),
                let msg = String(data: msgBytes, encoding: .utf8)
            else {
                return Just(ClientAction.clientError(.authChallengeError(1)))
                    .eraseToAnyPublisher()
            }

            let dict = msg.split(separator: ",")
                .compactMap { part -> (String, String)? in
                    guard part.count > 2 else { return nil }
                    let line1 = part.prefix(1)
                    let line2 = part.dropFirst(2)
                    return (String(line1), String(line2))
                }
                .reduce(into: [String: String]()) {
                    $0[$1.0] = $1.1
                }

            guard
                let serverNonce = dict["r"],
                let salt = dict["s"],
                let itrs = dict["i"],
                let iterations = Int(itrs)
            else {
                return Just(ClientAction.clientError(.authChallengeError(2)))
                    .eraseToAnyPublisher()
            }

            guard serverNonce.hasPrefix(clientNonce) else {
                return Just(ClientAction.clientError(.authChallengeError(3)))
                    .eraseToAnyPublisher()
            }

            guard
                let passData = state.authState.pass?.data(using: .utf8),
                let saltData = Data(base64Encoded: salt, options: Data.Base64DecodingOptions(rawValue: 0)),
                let clientKeyData = "Client Key".data(using: .utf8),
                let serverKeyData = "Server Key".data(using: .utf8)
            else {
                return Just(ClientAction.clientError(.authChallengeError(4)))
                    .eraseToAnyPublisher()
            }

            let saltedPass = passData.pbkdf2(salt: saltData, rounds: iterations)
            let clientFinalMessageBare = "c=biws,r=\(serverNonce)"
            let clientKey = clientKeyData.hmac(key: saltedPass)
            let storedKey = clientKey.sha1()
            let authMessage = "n=\(state.jid.localPart),r=\(clientNonce),\(msg),\(clientFinalMessageBare)"
            let clientSignature = authMessage.data(using: .utf8)?.hmac(key: storedKey) ?? Data()
            let clientProof = clientKey.xor(other: clientSignature)
            let serverKey = serverKeyData.hmac(key: saltedPass)
            let serverSignature = authMessage.data(using: .utf8)?.hmac(key: serverKey)
            let clientFinalMessage = "\(clientFinalMessageBare),p=\(clientProof.base64EncodedString())"

            let req = XMLElement(
                name: "response",
                xmlns: "urn:ietf:params:xml:ns:xmpp-sasl",
                attributes: [:],
                content: clientFinalMessage.base64Encoded,
                nodes: []
            )
            return Just(ClientAction.streamAction(.streamSend(req)))
                .eraseToAnyPublisher()

        default:
            return Empty().eraseToAnyPublisher()
        }
    }
}

// MARK: For PLAIN
private extension AuthMiddleware {
    func makePlainChallenge(jid: JID, pass: String) -> String {
        let lreq = "\0\(jid.localPart)\0\(pass)"
        let utf8str = lreq.data(using: .utf8)
        let base64 = utf8str?.base64EncodedString(options: NSData.Base64EncodingOptions(rawValue: 0))
        return base64 ?? ""
    }
}

// MARK: For SCRAM-SHA-1
private extension AuthMiddleware {
    func scramSha1InitRequestString(jid: JID) -> (String, String) {
        let randString = randomString(length: 20)
        let lreq = "n,,n=\(jid.localPart),r=\(randString)"
        let utf8str = lreq.data(using: .utf8)
        let base64 = utf8str?.base64EncodedString(options: NSData.Base64EncodingOptions(rawValue: 0))
        return (base64 ?? "", randString)
    }

    func randomString(length: Int) -> String {
        let letters = AuthMiddleware.ALPHABET
        return String((0 ..< length).map { _ in letters.randomElement() ?? "A" })
    }
}
