import Foundation
import CryptoKit

/// Drives a Noise handshake to completion, then yields the pair of transport `CipherState`s.
/// Supports exactly the patterns the sync layer uses: `XXpsk0` (pairing, mutual auth bound
/// to the QR pairing secret) and `IK` (reconnect, initiator already pins the host static).
struct NoiseHandshake {
    enum Pattern {
        case xxpsk0   // -> psk, e / <- e, ee, s, es / -> s, se
        case ik       // <- s (pre) ; -> e, es, s, ss / <- e, ee, se

        var protocolName: String {
            switch self {
            case .xxpsk0: "Noise_XXpsk0_25519_ChaChaPoly_SHA256"
            case .ik: "Noise_IK_25519_ChaChaPoly_SHA256"
            }
        }
        var usesPSK: Bool { self == .xxpsk0 }
        var messagePatterns: [[String]] {
            switch self {
            case .xxpsk0: [["psk", "e"], ["e", "ee", "s", "es"], ["s", "se"]]
            case .ik: [["e", "es", "s", "ss"], ["e", "ee", "se"]]
            }
        }
    }

    private var symmetric: SymmetricState
    private let initiator: Bool
    private let pattern: Pattern
    private let psk: Data?

    // Local/remote keys.
    private let s: Curve25519.KeyAgreement.PrivateKey   // local static
    private var e: Curve25519.KeyAgreement.PrivateKey?  // local ephemeral
    private var rs: Data?                                // remote static (public)
    private var re: Data?                                // remote ephemeral (public)

    private var patternIndex = 0

    /// - Parameters:
    ///   - remoteStatic: required for IK initiator (pins the host); the responder learns the
    ///     peer static during the handshake.
    ///   - psk: the 32-byte pairing secret for `xxpsk0`.
    init(
        pattern: Pattern, initiator: Bool,
        staticKey: Curve25519.KeyAgreement.PrivateKey,
        remoteStatic: Data? = nil, psk: Data? = nil, prologue: Data = Data()
    ) {
        self.pattern = pattern
        self.initiator = initiator
        self.s = staticKey
        self.rs = remoteStatic
        self.psk = psk
        self.symmetric = SymmetricState(protocolName: pattern.protocolName)
        symmetric.mixHash(prologue)
        applyPreMessages()
    }

    // IK pre-message "<- s": both sides MixHash the responder's static public key.
    private mutating func applyPreMessages() {
        guard pattern == .ik else { return }
        if initiator {
            if let rs { symmetric.mixHash(rs) }
        } else {
            symmetric.mixHash(s.publicKey.rawRepresentation)
        }
    }

    var isComplete: Bool { patternIndex >= pattern.messagePatterns.count }

    /// True if it's this side's turn to call `writeMessage`.
    var isMyTurnToWrite: Bool {
        guard !isComplete else { return false }
        // Initiator writes even-indexed messages, responder odd-indexed.
        return (patternIndex % 2 == 0) == initiator
    }

    mutating func writeMessage(payload: Data = Data()) throws -> Data {
        let tokens = pattern.messagePatterns[patternIndex]
        var buffer = Data()
        for token in tokens {
            switch token {
            case "psk":
                guard let psk else { throw NoiseError.unexpectedToken("psk without key") }
                symmetric.mixKeyAndHash(psk)
            case "e":
                let ek = NoiseDH.generate()
                e = ek
                let pub = ek.publicKey.rawRepresentation
                buffer.append(pub)
                symmetric.mixHash(pub)
                if pattern.usesPSK { symmetric.mixKey(pub) }
            case "s":
                let enc = try symmetric.encryptAndHash(s.publicKey.rawRepresentation)
                buffer.append(enc)
            case "ee", "es", "se", "ss":
                try mixDH(token)
            default:
                throw NoiseError.unexpectedToken(token)
            }
        }
        buffer.append(try symmetric.encryptAndHash(payload))
        patternIndex += 1
        return buffer
    }

    mutating func readMessage(_ message: Data) throws -> Data {
        let tokens = pattern.messagePatterns[patternIndex]
        var cursor = message.startIndex
        func take(_ n: Int) throws -> Data {
            guard message.distance(from: cursor, to: message.endIndex) >= n else {
                throw NoiseError.messageTooShort
            }
            let end = message.index(cursor, offsetBy: n)
            defer { cursor = end }
            return Data(message[cursor..<end])
        }

        for token in tokens {
            switch token {
            case "psk":
                guard let psk else { throw NoiseError.unexpectedToken("psk without key") }
                symmetric.mixKeyAndHash(psk)
            case "e":
                let pub = try take(32)
                re = pub
                symmetric.mixHash(pub)
                if pattern.usesPSK { symmetric.mixKey(pub) }
            case "s":
                let len = symmetric.cipher.hasKey ? 48 : 32  // +16 AEAD tag if encrypting
                let blob = try take(len)
                rs = try symmetric.decryptAndHash(blob)
            case "ee", "es", "se", "ss":
                try mixDH(token)
            default:
                throw NoiseError.unexpectedToken(token)
            }
        }
        let rest = Data(message[cursor..<message.endIndex])
        let payload = try symmetric.decryptAndHash(rest)
        patternIndex += 1
        return payload
    }

    private mutating func mixDH(_ token: String) throws {
        let ikm: Data
        switch token {
        case "ee":
            ikm = try NoiseDH.agree(req(e), req(re))
        case "es":
            ikm = initiator ? try NoiseDH.agree(req(e), req(rs))
                            : try NoiseDH.agree(s, req(re))
        case "se":
            ikm = initiator ? try NoiseDH.agree(s, req(re))
                            : try NoiseDH.agree(req(e), req(rs))
        case "ss":
            ikm = try NoiseDH.agree(s, req(rs))
        default:
            throw NoiseError.unexpectedToken(token)
        }
        symmetric.mixKey(ikm)
    }

    private func req(_ key: Curve25519.KeyAgreement.PrivateKey?) throws -> Curve25519.KeyAgreement.PrivateKey {
        guard let key else { throw NoiseError.handshakeNotFinished }
        return key
    }
    private func req(_ data: Data?) throws -> Data {
        guard let data else { throw NoiseError.handshakeNotFinished }
        return data
    }

    /// The peer's static public key, learned during the handshake (used to pin/identify).
    var remoteStaticKey: Data? { rs }

    /// After the final message, derive the two transport ciphers. Initiator: (send, recv);
    /// responder: (recv, send) — so we return them already oriented for this side.
    func transportCiphers() throws -> (send: CipherState, recv: CipherState) {
        guard isComplete else { throw NoiseError.handshakeNotFinished }
        let (c1, c2) = symmetric.split()
        return initiator ? (send: c1, recv: c2) : (send: c2, recv: c1)
    }

    /// Handshake hash — a unique channel binding usable for verification codes.
    var handshakeHash: Data { symmetric.h }
}
