import Foundation
import NucleicProtocol

/// A session-transfer failure, source side (mesh P5).
public enum SessionTransferError: Error, Sendable, Equatable {
    /// The destination rejected the offer (or later staging) with a typed reason.
    case rejected(TransferRejectReason, message: String)
    /// The channel dropped mid-transfer.
    case disconnected(String)
    /// The destination sent an out-of-sequence reply.
    case unexpectedReply
    /// This session already has an active transfer.
    case alreadyTransferring(SessionID)
    /// The transfer completed and the source session was tombstoned, but the destination never
    /// acked the commit before the channel dropped. The session is safely moved (archived +
    /// moved-to marker) and NOT running here; a reconnect re-sends the idempotent commit.
    case committedButUnconfirmed
}

/// Drives a session transfer from the **source** Mac as a 2-phase commit that never lets the
/// session run in two places (mesh P5): offer → accept → stream chunks → the destination stages
/// non-runnable and replies `transferReady` → the source **durably tombstones** → commit → the
/// destination activates and acks → the source discards the worktree.
///
/// The tombstone happens strictly *before* the commit is sent, so a crash/disconnect anywhere
/// leaves at most one runnable copy: before tombstone the session still lives here (the staged
/// copy is inert); after tombstone the source is archived+moved and only a committed destination
/// runs it.
public actor SessionTransferCoordinator {
    public struct Config: Sendable {
        /// Max chunks in flight before waiting for acks (stop-and-wait window).
        public var window: Int
        public init(window: Int = 4) { self.window = window }
    }

    private let session: Session
    private let project: Project
    private let worktree: Worktree
    private let destinationDeviceID: String
    private let channel: any TransferChannel
    private let store: any SessionMetadataStore
    private let worktrees: any WorktreeManaging
    private let config: Config
    private let package: SessionCarbonCopy

    // Reply FIFO — a detached pump drains the channel's replies; the linear flow pulls them.
    private var pending: [TransferReply] = []
    private var waiter: CheckedContinuation<TransferReply?, Never>?
    private var streamClosed = false
    private var pumpStarted = false

    public init(
        session: Session, project: Project, worktree: Worktree,
        destinationDeviceID: String, package: SessionCarbonCopy,
        channel: any TransferChannel, store: any SessionMetadataStore,
        worktrees: any WorktreeManaging, config: Config = Config()
    ) {
        self.session = session
        self.project = project
        self.worktree = worktree
        self.destinationDeviceID = destinationDeviceID
        self.package = package
        self.channel = channel
        self.store = store
        self.worktrees = worktrees
        self.config = config
    }

    /// Run the transfer to completion. Throws `SessionTransferError` on any failure; on success
    /// the source session is tombstoned and its worktree/branch discarded.
    public func run(now: Date = Date()) async throws {
        startPump()
        // Claim the lock — throws if this session is already transferring.
        let lock = SessionTransferLock(
            transferID: package.transferID, sessionID: session.id, direction: .outbound,
            peerDeviceID: destinationDeviceID, state: .offering, createdAt: now, updatedAt: now)
        do {
            try await store.beginTransfer(lock)
        } catch is SessionTransferConflict {
            throw SessionTransferError.alreadyTransferring(session.id)
        }

        do {
            try await drive(now: now)
        } catch {
            // Abandon: if we never tombstoned, the session still lives here — clear the lock and
            // staging so the user can retry. (A tombstoned failure is rethrown as-is; the session
            // is already moved and the lock stays at `.tombstoned` for reconnect recovery.)
            if case SessionTransferError.committedButUnconfirmed = error {
                throw error
            }
            try? await store.clearTransfer(transferID: package.transferID)
            package.cleanup()
            throw error
        }
    }

    private func drive(now: Date) async throws {
        // 1. Offer.
        await channel.send(.transferOffer(package.offer))
        let accept = try await expectAccept()

        // 2. Build the bundle against what the destination has, then stream every item — except
        //    the ones the destination already holds whole (`alreadyHaveItems`, e.g. a transcript
        //    it adopted from its mesh mirror): those it staged itself, so streaming them would
        //    only re-send bytes the mesh already synced.
        _ = try await package.buildBundle(haveSHAs: accept.haveSHAs)
        try await store.updateTransferState(transferID: package.transferID, to: .streaming, at: now)
        for kind in package.streamOrder where !accept.alreadyHaveItems.contains(kind) {
            try await streamItem(kind, resumeFrom: accept.resumeFrom[kind] ?? 0)
        }

        // 3. Await the destination staging the session non-runnable.
        try await store.updateTransferState(transferID: package.transferID, to: .awaitingReady, at: now)
        try await expectReady()

        // 4. THE COMMIT POINT — durably tombstone the source before asking the destination to
        //    activate. Now at most the destination can run this session.
        try await store.tombstoneSession(id: session.id, movedToDeviceID: destinationDeviceID, at: now)
        try await store.updateTransferState(transferID: package.transferID, to: .tombstoned, at: now)

        // 5. Commit (idempotent) and await the destination's activation ack. Past the tombstone
        //    the source is durably moved: whatever happens next — a bad reply, a dropped channel,
        //    even a local DB write failing — we must NOT revive the session here. So the ENTIRE
        //    post-tombstone tail is wrapped: any throw surfaces as `committedButUnconfirmed`, which
        //    the caller treats as "moved, awaiting confirmation" and which leaves the lock at
        //    `.tombstoned` for reconnect recovery — never clearing it or re-listing the source.
        do {
            await channel.send(.transferCommit(package.transferID))
            try await expectCommitted()
            try await store.updateTransferState(transferID: package.transferID, to: .committed, at: now)
        } catch {
            throw SessionTransferError.committedButUnconfirmed
        }

        // 6. Cleanup — the branch's work lives on the destination now.
        try? await worktrees.discardAfterTransfer(worktree, in: project)
        try? await store.clearTransfer(transferID: package.transferID)
        package.cleanup()
    }

    /// Cancel an in-flight transfer before the commit point.
    public func cancel() async {
        await channel.send(.transferCancel(package.transferID))
        try? await store.clearTransfer(transferID: package.transferID)
        package.cleanup()
    }

    /// Re-drive the commit tail of a durably-`tombstoned` transfer after a relaunch (mesh P5
    /// recovery). The source already moved (archived + moved-to marker) but never heard the
    /// destination's committed ack — so we re-send the idempotent `transferCommit` over a
    /// re-established channel and, if the destination confirms, discard the now-stale source
    /// worktree and clear the lock. Returns whether the transfer is now fully committed.
    ///
    /// Best-effort and safe to call repeatedly: a destination that already activated re-acks; one
    /// that still has it staged activates now; an unreachable or since-relaunched destination
    /// (which lost its in-memory staging) leaves the lock at `.tombstoned` for a later attempt —
    /// the source is never revived either way. A bounded wait keeps a connected-but-silent peer
    /// from wedging the recovery task.
    public static func recoverTombstoned(
        lock: SessionTransferLock, session: Session?, project: Project?,
        channel: any TransferChannel, store: any SessionMetadataStore,
        worktrees: any WorktreeManaging, now: Date = Date(),
        timeout: Duration = .seconds(20)
    ) async -> Bool {
        await channel.send(.transferCommit(lock.transferID))
        let confirmed = await withTaskGroup(of: Bool.self) { group -> Bool in
            group.addTask {
                for await reply in channel.replies() {
                    switch reply {
                    case .committed(let id) where id == lock.transferID: return true
                    case .reject, .disconnected: return false
                    default: continue   // stray ack from an interrupted stream — keep waiting
                    }
                }
                return false
            }
            group.addTask {
                try? await Task.sleep(for: timeout)
                return false
            }
            let result = await group.next() ?? false
            group.cancelAll()
            return result
        }
        guard confirmed else { return false }
        // The destination owns the branch now — drop the source worktree, then clear the lock.
        if let project, let path = session?.worktreePath, let branch = session?.branch,
           let baseSHA = session?.baseSHA, let session {
            let worktree = Worktree(
                sessionID: session.id, path: path, branch: branch, baseSHA: baseSHA,
                createdAt: session.createdAt)
            try? await worktrees.discardAfterTransfer(worktree, in: project)
        }
        try? await store.updateTransferState(transferID: lock.transferID, to: .committed, at: now)
        try? await store.clearTransfer(transferID: lock.transferID)
        return true
    }

    // MARK: - Streaming

    private func streamItem(_ kind: TransferItemKind, resumeFrom: Int) async throws {
        guard let item = package.item(kind) else { return }
        let descriptor = item.descriptor
        let handle = try FileHandle(forReadingFrom: item.url)
        defer { try? handle.close() }
        if resumeFrom > 0 {
            try handle.seek(toOffset: UInt64(resumeFrom * TransferChunk.maxDataBytes))
        }

        var unacked = 0
        var index = resumeFrom
        while index < descriptor.chunkCount {
            let data = (try handle.read(upToCount: TransferChunk.maxDataBytes)) ?? Data()
            await channel.send(.transferChunk(TransferChunk(
                transferID: package.transferID, kind: kind, chunkIndex: index,
                chunkCount: descriptor.chunkCount, sha256: descriptor.sha256, data: data)))
            unacked += 1
            index += 1
            while unacked >= config.window {
                try await expectAck(kind: kind)
                unacked -= 1
            }
        }
        while unacked > 0 {
            try await expectAck(kind: kind)
            unacked -= 1
        }
    }

    // MARK: - Reply pump + typed expectations

    private func startPump() {
        guard !pumpStarted else { return }
        pumpStarted = true
        let stream = channel.replies()
        Task { [weak self] in
            for await reply in stream { await self?.enqueue(reply) }
            await self?.closeStream()
        }
    }

    private func enqueue(_ reply: TransferReply) {
        if let w = waiter { waiter = nil; w.resume(returning: reply) }
        else { pending.append(reply) }
    }
    private func closeStream() {
        streamClosed = true
        if let w = waiter { waiter = nil; w.resume(returning: nil) }
    }
    private func nextReply() async -> TransferReply? {
        if !pending.isEmpty { return pending.removeFirst() }
        if streamClosed { return nil }
        return await withCheckedContinuation { waiter = $0 }
    }

    private func expectAccept() async throws -> TransferAccept {
        switch await nextReply() {
        case .accept(let a): return a
        case .reject(let r): throw SessionTransferError.rejected(r.reason, message: r.message)
        case .disconnected(let why): throw SessionTransferError.disconnected(why)
        case .none: throw SessionTransferError.disconnected("closed")
        default: throw SessionTransferError.unexpectedReply
        }
    }

    private func expectAck(kind: TransferItemKind) async throws {
        switch await nextReply() {
        case .chunkAck: return
        case .reject(let r): throw SessionTransferError.rejected(r.reason, message: r.message)
        case .disconnected(let why): throw SessionTransferError.disconnected(why)
        case .none: throw SessionTransferError.disconnected("closed")
        default: throw SessionTransferError.unexpectedReply
        }
    }

    private func expectReady() async throws {
        switch await nextReply() {
        case .ready: return
        case .reject(let r): throw SessionTransferError.rejected(r.reason, message: r.message)
        case .disconnected(let why): throw SessionTransferError.disconnected(why)
        case .none: throw SessionTransferError.disconnected("closed")
        default: throw SessionTransferError.unexpectedReply
        }
    }

    private func expectCommitted() async throws {
        switch await nextReply() {
        case .committed: return
        case .reject(let r): throw SessionTransferError.rejected(r.reason, message: r.message)
        case .disconnected(let why): throw SessionTransferError.disconnected(why)
        case .none: throw SessionTransferError.disconnected("closed")
        default: throw SessionTransferError.unexpectedReply
        }
    }
}
