//
//  DPoP.swift
//  AT-IRC
//
//  ES256 DPoP key, JWK thumbprint, and proof generation (RFC 9449).
//

import Foundation
import CryptoKit

struct DPoPKey {
    let privateKey: P256.Signing.PrivateKey

    init() {
        privateKey = P256.Signing.PrivateKey()
    }

    init(rawRepresentation: Data) throws {
        privateKey = try P256.Signing.PrivateKey(rawRepresentation: rawRepresentation)
    }

    var rawRepresentation: Data {
        privateKey.rawRepresentation
    }

    /// Public key as an EC P-256 JWK.
    var publicJWK: [String: String] {
        let x963 = privateKey.publicKey.x963Representation
        let x = x963.subdata(in: 1..<33)
        let y = x963.subdata(in: 33..<65)
        return [
            "crv": "P-256",
            "kty": "EC",
            "x": Base64URL.encode(x),
            "y": Base64URL.encode(y),
        ]
    }

    /// RFC 7638 JWK thumbprint (SHA-256 over the canonical JWK, base64url).
    var thumbprint: String {
        let jwk = publicJWK
        let canonical = "{\"crv\":\"\(jwk["crv"]!)\",\"kty\":\"\(jwk["kty"]!)\",\"x\":\"\(jwk["x"]!)\",\"y\":\"\(jwk["y"]!)\"}"
        let digest = SHA256.hash(data: Data(canonical.utf8))
        return Base64URL.encode(Data(digest))
    }

    /// Build a signed DPoP proof JWT for one request.
    /// - Parameters:
    ///   - method: HTTP method (`GET`, `POST`, ...).
    ///   - url: Request URL; query and fragment are stripped for `htu`.
    ///   - nonce: Server-provided `DPoP-Nonce`, when available.
    ///   - accessToken: When calling a resource server, the token the proof is bound to (`ath`).
    func proof(method: String, url: URL, nonce: String? = nil, accessToken: String? = nil) throws -> String {
        let header: [String: Any] = [
            "alg": "ES256",
            "typ": "dpop+jwt",
            "jwk": publicJWK,
        ]
        var claims: [String: Any] = [
            "jti": UUID().uuidString,
            "htm": method.uppercased(),
            "htu": Self.htu(url),
            "iat": Int(Date().timeIntervalSince1970),
        ]
        if let nonce, !nonce.isEmpty {
            claims["nonce"] = nonce
        }
        if let accessToken {
            let digest = SHA256.hash(data: Data(accessToken.utf8))
            claims["ath"] = Base64URL.encode(Data(digest))
        }

        let headerData = try JSONSerialization.data(withJSONObject: header, options: [.sortedKeys])
        let claimsData = try JSONSerialization.data(withJSONObject: claims, options: [.sortedKeys])
        let signingInput = Base64URL.encode(headerData) + "." + Base64URL.encode(claimsData)
        let signature = try privateKey.signature(for: Data(signingInput.utf8))
        return signingInput + "." + Base64URL.encode(signature.rawRepresentation)
    }

    /// `htu` is the request URI without query or fragment (RFC 9449 §4.2).
    static func htu(_ url: URL) -> String {
        guard var components = URLComponents(url: url, resolvingAgainstBaseURL: false) else {
            return url.absoluteString
        }
        components.query = nil
        components.fragment = nil
        return components.string ?? url.absoluteString
    }
}

enum DPoPError: LocalizedError {
    case signingFailed
    var errorDescription: String? {
        switch self {
        case .signingFailed: "Could not sign DPoP proof"
        }
    }
}