import Foundation
import NucleicProtocol

/// Keeps this member's Carbon store caught up with the mesh (docs/CARBON_SHARDING.md §8.2,
/// §8.5): on a pushed `carbonHeadAdvanced` — or a reconcile pass against a peer's full head
/// state — it walks the manifest chain back to something it holds, fetches the missing
/// shards, and only then advances its own head, so a gossiped head is always an honest
/// "I hold everything reachable" claim (the durability ack, §5.3).
///
/// Backfill (a new replica member, a returning device, a replaced device) is the same code
/// with no special mode: exchange heads → walk chains → fetch missing.
///
/// Transport is injected (the `TranscriptSyncCoordinator` pattern) so the whole replication
/// discipline — including 3-member convergence through partition/rejoin — is testable against
/// in-memory fakes; the `PeerClient` adapter binds the closures to real wire verbs.
public actor CarbonSyncCoordinator {
    /// How a peer is currently reachable — drives the §8.4 relay payload policy.
    public enum TransportClass: Sendable, Equatable {
        case direct  // LAN / tailnet: no size limit
        case relay   // metered fallback: steady-state manifests only, ≤ the payload limit
    }

    /// Wire access, abstracted. Every closure returns what the peer *served* (nil/empty on
    /// unreachable or refused) — the coordinator owns fallthrough across holders.
    public struct PeerTransport: Sendable {
        public let fetchHeads: @Sendable (_ peer: String, _ projectUUID: UUID?) async -> [CarbonHead]?
        public let fetchManifests:
            @Sendable (_ peer: String, _ projectUUID: UUID, _ ids: [String]) async -> [String: Data]
        public let fetchShards:
            @Sendable (_ peer: String, _ projectUUID: UUID, _ ids: [String]) async -> [String: Data]
        public let transportClass: @Sendable (_ peer: String) async -> TransportClass

        public init(
            fetchHeads: @escaping @Sendable (String, UUID?) async -> [CarbonHead]?,
            fetchManifests: @escaping @Sendable (String, UUID, [String]) async -> [String: Data],
            fetchShards: @escaping @Sendable (String, UUID, [String]) async -> [String: Data],
            transportClass: @escaping @Sendable (String) async -> TransportClass = { _ in .direct }
        ) {
            self.fetchHeads = fetchHeads
            self.fetchManifests = fetchManifests
            self.fetchShards = fetchShards
            self.transportClass = transportClass
        }
    }

    /// D5: the relay carries a chain advance only when the manifest's total shard bytes fit
    /// under this; larger payloads defer until a direct path exists (§8.4).
    public static let defaultRelayPayloadLimit: UInt64 = 4 * 1024 * 1024

    private let store: CarbonShardStore
    private let keyProvider: CarbonShardStore.KeyProvider
    private let transport: PeerTransport
    private let maxConcurrentChains: Int
    private let relayPayloadLimit: UInt64

    /// One in-flight sync per stream key (reserve synchronously before any await).
    private var inFlight: Set<CarbonStreamKey> = []
    private var activeCount = 0
    /// The newest head each peer has gossiped/served per stream — the §5.3 durability input.
    private var peerHeads: [String: [CarbonStreamKey: CarbonHead]] = [:]
    /// Chains whose payload exceeded the relay limit while only relay peers could serve them —
    /// surfaced as `reduced redundancy (awaiting direct link)` and retried on reconcile.
    private var deferredAwaitingDirect: Set<CarbonStreamKey> = []

    public init(
        store: CarbonShardStore, keyProvider: @escaping CarbonShardStore.KeyProvider,
        transport: PeerTransport, maxConcurrentChains: Int = 3,
        relayPayloadLimit: UInt64 = CarbonSyncCoordinator.defaultRelayPayloadLimit
    ) {
        self.store = store
        self.keyProvider = keyProvider
        self.transport = transport
        self.maxConcurrentChains = maxConcurrentChains
        self.relayPayloadLimit = relayPayloadLimit
    }

    // MARK: - Entry points

    /// A peer pushed `carbonHeadAdvanced` (§8.2 step 1) — or its head arrived via gossip.
    /// `candidates` are the other reachable peers to fall through to (owner-first ordering is
    /// the caller's, as with transcript sync).
    public func observeHeadAdvanced(
        _ head: CarbonHead, from peer: String, candidates: [String] = []
    ) async {
        recordPeerHead(head, peer: peer)
        await syncChainIfNeeded(to: head, holders: [peer] + candidates.filter { $0 != peer })
    }

    /// On (re)connect or a reconcile tick: exchange full head state with `peers` (heads are
    /// state, not events — missed pushes are irrelevant, §8.2 step 4) and catch up every
    /// chain that is ahead of ours. Also retries relay-deferred chains.
    public func reconcile(projectUUID: UUID?, peers: [String]) async {
        for peer in peers {
            guard let heads = await transport.fetchHeads(peer, projectUUID) else { continue }
            for head in heads {
                recordPeerHead(head, peer: peer)
                await syncChainIfNeeded(to: head, holders: [peer] + peers.filter { $0 != peer })
            }
        }
    }

    // MARK: - Durability accounting (§5.3)

    public enum Durability: Sendable, Equatable {
        case protected(holders: Int)
        case reducedRedundancy(reason: String)
        case notProtected(reason: String)
    }

    /// How durable one of OUR heads is: which replica-set members have gossiped a head at or
    /// past it (an honest claim — they hold everything reachable). Mesh-durable when
    /// `count >= min(2, |replicaSet|)`, degrading per the D3 table, never blocking anything.
    public func durability(of head: CarbonHead, replicaSet: Set<String>) -> Durability {
        guard !replicaSet.isEmpty else {
            return .notProtected(reason: "no eligible replica members")
        }
        let ackCount = replicaSet.filter { member in
            guard let peerHead = peerHeads[member]?[head.key] else { return false }
            return peerHead.ownerGeneration > head.ownerGeneration
                || (peerHead.ownerGeneration == head.ownerGeneration
                    && peerHead.snapshotSeq >= head.snapshotSeq)
        }.count
        if deferredAwaitingDirect.contains(head.key) && ackCount < 2 {
            return .reducedRedundancy(reason: "awaiting direct link")
        }
        if ackCount >= min(2, replicaSet.count) {
            return .protected(holders: ackCount + 1)  // + the owner itself
        }
        if ackCount == 1 { return .reducedRedundancy(reason: "one replica holds it") }
        return .notProtected(reason: "no replica holds it yet")
    }

    /// Chains currently deferred because only relay-class peers can serve an over-limit
    /// payload — for the durability chip.
    public func chainsAwaitingDirectLink() -> Set<CarbonStreamKey> { deferredAwaitingDirect }

    /// What gossip says each peer holds (§8.2: heads are *state*) — the same input `durability`
    /// scores, exposed read-only for the Carbon Viewer's per-member rows. A peer absent here has
    /// simply never gossiped to us; that is "nothing reported", not "holds nothing".
    public func peerHeadState() -> [String: [CarbonStreamKey: CarbonHead]] { peerHeads }

    // MARK: - Chain sync

    private func recordPeerHead(_ head: CarbonHead, peer: String) {
        let current = peerHeads[peer]?[head.key]
        let advances = current.map {
            head.ownerGeneration > $0.ownerGeneration
                || (head.ownerGeneration == $0.ownerGeneration && head.snapshotSeq > $0.snapshotSeq)
        } ?? true
        if advances { peerHeads[peer, default: [:]][head.key] = head }
    }

    private func syncChainIfNeeded(to target: CarbonHead, holders: [String]) async {
        let key = target.key
        // §10 fencing at the door: a target at or below our head is nothing to do; a
        // stale-generation target must never advance us.
        if let local = await store.heads(projectUUID: key.projectUUID)[key] {
            let ahead =
                target.ownerGeneration > local.ownerGeneration
                || (target.ownerGeneration == local.ownerGeneration
                    && target.snapshotSeq > local.snapshotSeq)
            guard ahead else { return }
        }
        // Mirroring requires the project key: chain walking decrypts manifests to find
        // parents (§8.2 step 2). A keyless member reports `keyMissing` via gossip instead.
        guard keyProvider(key.projectUUID) != nil else { return }
        guard activeCount < maxConcurrentChains, !inFlight.contains(key) else { return }
        inFlight.insert(key)
        activeCount += 1
        await runChainSync(target: target, holders: holders)
        inFlight.remove(key)
        activeCount -= 1
    }

    private func runChainSync(target: CarbonHead, holders: [String]) async {
        let projectUUID = target.key.projectUUID
        guard let key = keyProvider(projectUUID) else { return }
        let crypto = CarbonCrypto(projectKey: key)

        // Walk parentManifest links from the target until we reach a manifest we hold (or the
        // chain start), collecting the missing segment newest→oldest.
        var segment: [(id: String, manifest: CarbonManifest)] = []
        var cursor: String? = target.manifestID
        while let id = cursor {
            if await store.manifestBlob(id: id, projectUUID: projectUUID) != nil { break }
            guard let blob = await fetchManifestBlob(
                id: id, projectUUID: projectUUID, holders: holders),
                let manifest = try? CarbonManifest.open(blob, with: crypto)
            else { return }  // chain incomplete — retry on a later pass/reconcile
            // Persist the manifest blob now (verified put); the HEAD advances only after the
            // whole segment's shards are held (§8.2 step 3).
            guard (try? await store.putManifest(
                ciphertext: blob, expectedID: id, projectUUID: projectUUID)) != nil
            else { return }
            segment.append((id, manifest))
            cursor = manifest.parentManifest
        }

        // §8.4 relay policy: if every serving peer is relay-class and any manifest's payload
        // exceeds the limit, defer the chain until a direct path exists.
        let oversize = segment.contains {
            $0.manifest.shards.reduce(UInt64(0)) { $0 + $1.byteCount } > relayPayloadLimit
        }
        if oversize {
            var anyDirect = false
            for holder in holders {
                if await transport.transportClass(holder) == .direct {
                    anyDirect = true
                    break
                }
            }
            if !anyDirect {
                deferredAwaitingDirect.insert(target.key)
                return
            }
        }
        deferredAwaitingDirect.remove(target.key)

        // Fetch missing shards oldest→newest, then advance the head last (honest-head rule).
        for (_, manifest) in segment.reversed() {
            let missing = await store.missingShards(for: manifest)
            guard await fetchShards(missing, projectUUID: projectUUID, holders: holders) else {
                return
            }
        }
        try? await store.advanceHead(target)
    }

    private func fetchManifestBlob(
        id: String, projectUUID: UUID, holders: [String]
    ) async -> Data? {
        for holder in holders {
            if let blob = await transport.fetchManifests(holder, projectUUID, [id])[id] {
                return blob
            }
        }
        return nil
    }

    /// Fetch `ids` with fallthrough across holders (batched per holder; ≤256-id batching is
    /// the wire adapter's concern). Returns true iff every shard landed verified.
    private func fetchShards(
        _ ids: [String], projectUUID: UUID, holders: [String]
    ) async -> Bool {
        var remaining = ids
        for holder in holders {
            guard !remaining.isEmpty else { break }
            let served = await transport.fetchShards(holder, projectUUID, remaining)
            var still: [String] = []
            for id in remaining {
                if let blob = served[id],
                    (try? await store.putShard(
                        ciphertext: blob, expectedID: id, projectUUID: projectUUID)) != nil
                {
                    continue
                }
                still.append(id)
            }
            remaining = still
        }
        return remaining.isEmpty
    }
}
