// IdentityManager.swift

import Foundation
import CryptoKit
import Security
import os.log

private let log = Logger(subsystem: "com.btmessage.app", category: "Identity")

public class IdentityManager: ObservableObject {
    public static let shared = IdentityManager()

    @Published public private(set) var localPeerID: String
    @Published public private(set) var publicKey: HybridPublicKey
    @Published public private(set) var error: IdentityError?

    private var privateKey: HybridPrivateKey

    private init() {
        if let (peerID, pub, priv) = IdentityManager.loadIdentity() {
            self.localPeerID = peerID
            self.publicKey = pub
            self.privateKey = priv
            log.info("Identity loaded: \(peerID.prefix(8))...")
        } else {
            log.info("Creating new identity...")
            do {
                let peerID = UUID().uuidString
                let (pub, priv) = try HybridCrypto.generateKeyPair(peerID: peerID)
                self.localPeerID = peerID
                self.publicKey = pub
                self.privateKey = priv
                IdentityManager.saveIdentity(peerID: peerID, privKey: priv.xwing)
                log.info("New identity created: \(peerID.prefix(8))...")
            } catch {
                log.fault("CRITICAL: Failed to generate identity: \(error)")
                fatalError("Cannot create identity: \(error)")
            }
        }
    }

    public func encrypt(message: Data, for recipientPublicKey: HybridPublicKey) throws -> EncryptedMessage {
        try HybridCrypto.encrypt(message: message, recipientPublicKey: recipientPublicKey)
    }

    public func decrypt(message: EncryptedMessage) throws -> Data {
        try HybridCrypto.decrypt(message: message, recipientPrivateKey: privateKey)
    }

    // MARK: - Keychain

    private static let peerIDKey = "btmessage.identity.peerID"
    private static let xwingKey = "btmessage.identity.xwing"

    private static func loadIdentity() -> (String, HybridPublicKey, HybridPrivateKey)? {
        guard
            let peerIDData = keychainLoad(key: peerIDKey),
            let peerID = String(data: peerIDData, encoding: .utf8),
            let keyData = keychainLoad(key: xwingKey)
        else {
            log.debug("No existing identity found in keychain")
            return nil
        }

        do {
            let xwing = try XWingMLKEM768X25519.PrivateKey(integrityCheckedRepresentation: keyData)
            let pub = HybridPublicKey(rawBytes: xwing.publicKey.rawRepresentation, peerID: peerID)
            let priv = HybridPrivateKey(xwing: xwing)
            return (peerID, pub, priv)
        } catch {
            log.error("Failed to load identity: \(error)")
            return nil
        }
    }

    private static func saveIdentity(peerID: String, privKey: XWingMLKEM768X25519.PrivateKey) {
        let keyData = Data(privKey.integrityCheckedRepresentation)

        let peerSaved = keychainSave(key: peerIDKey, data: peerID.data(using: .utf8)!)
        let keySaved = keychainSave(key: xwingKey, data: keyData)

        if peerSaved && keySaved {
            log.info("Identity saved to keychain")
        } else {
            log.error("Failed to save identity - peer:\(peerSaved) key:\(keySaved)")
        }
    }

    @discardableResult
    private static func keychainSave(key: String, data: Data) -> Bool {
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: key,
            kSecValueData as String: data,
            kSecAttrAccessible as String: kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly
        ]
        SecItemDelete(query as CFDictionary)
        let status = SecItemAdd(query as CFDictionary, nil)
        return status == errSecSuccess
    }

    private static func keychainLoad(key: String) -> Data? {
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: key,
            kSecReturnData as String: true,
            kSecMatchLimit as String: kSecMatchLimitOne
        ]
        var item: CFTypeRef?
        guard SecItemCopyMatching(query as CFDictionary, &item) == errSecSuccess else { return nil }
        return item as? Data
    }
}

public enum IdentityError: Error, LocalizedError {
    case keyGenerationFailed(Error)
    case keychainSaveFailed
    case keychainLoadFailed

    public var errorDescription: String? {
        switch self {
        case .keyGenerationFailed(let err): return "Key generation failed: \(err.localizedDescription)"
        case .keychainSaveFailed: return "Failed to save to keychain"
        case .keychainLoadFailed: return "Failed to load from keychain"
        }
    }
}