montana/Montana-iOS/Montana Messenger/Montana Messenger/Stores.swift

import Foundation
import Combine
import SwiftUI

struct Contact: Identifiable, Codable, Hashable {
    let accountID: String
    var name: String
    var xPubB64: String
    var addedAt: Date
    var id: String { accountID }
}

struct LocalMessage: Identifiable, Codable, Hashable {
    let id: Int
    let peer: String
    let outgoing: Bool
    let text: String
    let sentAt: Int
    var read: Bool
}

@MainActor
final class ContactsStore: ObservableObject {
    static let shared = ContactsStore()

    @Published private(set) var contacts: [Contact] = []

    private let key = "montana.contacts.v1"

    init() { load() }

    func load() {
        guard let data = UserDefaults.standard.data(forKey: key),
              let arr = try? JSONDecoder().decode([Contact].self, from: data) else {
            contacts = []
            return
        }
        contacts = arr.sorted { $0.name.localizedCaseInsensitiveCompare($1.name) == .orderedAscending }
    }

    private func persist() {
        if let data = try? JSONEncoder().encode(contacts) {
            UserDefaults.standard.set(data, forKey: key)
        }
    }

    func contact(byID id: String) -> Contact? {
        contacts.first { $0.accountID == id }
    }

    @discardableResult
    func addOrUpdate(_ c: Contact) -> Contact {
        if let idx = contacts.firstIndex(where: { $0.accountID == c.accountID }) {
            contacts[idx] = c
        } else {
            contacts.append(c)
        }
        contacts.sort { $0.name.localizedCaseInsensitiveCompare($1.name) == .orderedAscending }
        persist()
        return c
    }

    func delete(_ id: String) {
        contacts.removeAll { $0.accountID == id }
        persist()
    }

    func addByID(_ accountID: String, fallbackName: String? = nil) async throws -> Contact {
        let trimmed = accountID.lowercased().trimmingCharacters(in: .whitespacesAndNewlines)
        guard trimmed.count >= 4, trimmed.count <= 64,
              trimmed.allSatisfy({ "0123456789abcdef".contains($0) })
        else {
            throw NSError(domain: "Contacts", code: 1,
                          userInfo: [NSLocalizedDescriptionKey: "Невалидный Account ID"])
        }
        let acc = try await ApiClient.shared.getAccount(trimmed)
        let c = Contact(accountID: acc.account_id, name: acc.name,
                        xPubB64: acc.x_pub, addedAt: Date())
        return addOrUpdate(c)
    }
}

@MainActor
final class ChatsStore: ObservableObject {
    static let shared = ChatsStore()

    @Published private(set) var messagesByPeer: [String: [LocalMessage]] = [:]
    @Published private(set) var lastMessageID: Int = 0
    @Published private(set) var unreadByPeer: [String: Int] = [:]
    @Published private(set) var isOnline = false

    private let messagesKey = "montana.messages.v1"
    private let cursorKey  = "montana.cursor.v1"
    private let unreadKey  = "montana.unread.v1"

    private var pollTask: Task<Void, Never>?

    init() {
        if let d = UserDefaults.standard.data(forKey: messagesKey),
           let m = try? JSONDecoder().decode([String: [LocalMessage]].self, from: d) {
            messagesByPeer = m
        }
        lastMessageID = UserDefaults.standard.integer(forKey: cursorKey)
        if let d = UserDefaults.standard.data(forKey: unreadKey),
           let u = try? JSONDecoder().decode([String: Int].self, from: d) {
            unreadByPeer = u
        }
    }

    private func persist() {
        if let d = try? JSONEncoder().encode(messagesByPeer) {
            UserDefaults.standard.set(d, forKey: messagesKey)
        }
        UserDefaults.standard.set(lastMessageID, forKey: cursorKey)
        if let d = try? JSONEncoder().encode(unreadByPeer) {
            UserDefaults.standard.set(d, forKey: unreadKey)
        }
    }

    func messages(for peer: String) -> [LocalMessage] {
        messagesByPeer[peer] ?? []
    }

    func chats() -> [(peer: Contact, last: LocalMessage?)] {
        let all = ContactsStore.shared.contacts
        let withMsgs = all.compactMap { c -> (Contact, LocalMessage?)? in
            let msgs = messagesByPeer[c.accountID] ?? []
            guard let last = msgs.last else { return nil }
            return (c, last)
        }
        let withoutMsgs = all.filter { (messagesByPeer[$0.accountID] ?? []).isEmpty }
            .map { ($0, LocalMessage?.none) }
        return (withMsgs.sorted { ($0.1?.sentAt ?? 0) > ($1.1?.sentAt ?? 0) } + withoutMsgs)
    }

    func markRead(_ peer: String) {
        unreadByPeer[peer] = 0
        if let arr = messagesByPeer[peer] {
            messagesByPeer[peer] = arr.map { var m = $0; m.read = true; return m }
        }
        persist()
    }

    func absorb(envelope env: ApiEnvelope, identity: Identity) async {
        let me = identity.accountID
        let peer = (env.from == me) ? env.to : env.from
        let outgoing = env.from == me

        var contact = ContactsStore.shared.contact(byID: peer)
        if contact == nil {
            if let acc = try? await ApiClient.shared.getAccount(peer) {
                contact = ContactsStore.shared.addOrUpdate(
                    Contact(accountID: acc.account_id, name: acc.name,
                            xPubB64: acc.x_pub, addedAt: Date())
                )
            }
        }

        guard let xPub = contact?.xPubB64,
              let plaintext = MessageCrypto.decrypt(
                nonceB64: env.nonce, ciphertextB64: env.ciphertext,
                peerXPubB64: xPub, myAgreementKey: identity.agreementKey)
        else { return }

        var arr = messagesByPeer[peer] ?? []
        if arr.contains(where: { $0.id == env.id }) { return }
        arr.append(LocalMessage(id: env.id, peer: peer, outgoing: outgoing,
                                text: plaintext, sentAt: env.sent_at, read: outgoing))
        messagesByPeer[peer] = arr.sorted { $0.id < $1.id }

        if !outgoing {
            unreadByPeer[peer] = (unreadByPeer[peer] ?? 0) + 1
        }
        if env.id > lastMessageID { lastMessageID = env.id }
        persist()
    }

    func send(text: String, to peer: Contact, identity: Identity) async throws {
        let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { return }
        let (nonce, ct) = try MessageCrypto.encrypt(
            trimmed, recipientXPubB64: peer.xPubB64,
            mySigningKey: identity.signingKey, myAgreementKey: identity.agreementKey)
        let env = try await ApiClient.shared.send(
            to: peer.accountID, nonce: nonce, ciphertext: ct, identity: identity)
        await absorb(envelope: env, identity: identity)
    }

    func startPolling(identity: Identity) {
        pollTask?.cancel()
        let id = identity
        pollTask = Task { [weak self] in
            while !Task.isCancelled {
                guard let self else { return }
                do {
                    let envs = try await ApiClient.shared.inbox(since: self.lastMessageID, identity: id)
                    for env in envs {
                        await self.absorb(envelope: env, identity: id)
                    }
                    await MainActor.run { self.isOnline = true }
                } catch {
                    await MainActor.run { self.isOnline = false }
                }
                try? await Task.sleep(nanoseconds: 4_000_000_000)
            }
        }
    }

    func stopPolling() {
        pollTask?.cancel()
        pollTask = nil
        isOnline = false
    }

    func attachWebSocket(identity: Identity) {
        let id = identity
        WebSocketLink.shared.onEnvelope = { [weak self] env in
            guard let self else { return }
            Task { await self.absorb(envelope: env, identity: id) }
        }
        WebSocketLink.shared.connect(identity: id)
    }
}