import Foundation
import NucleicProtocol

/// The destination side of a session transfer (mesh P5): receives the offer, stages the session
/// **non-runnable** (git bundle imported, transcript written, native memory placed best-effort),
/// and — only on `transferCommit` — activates it as a live session. Idempotent on commit so a
/// re-delivered commit after a reconnect just re-acks.
///
/// One instance per host (owned by the app), keyed by `transferID`, so several transfers can
/// stage concurrently. All per-transfer mutable state lives in the actor.
public actor SessionTransferImporter {
    private let store: any SessionMetadataStore
    private let worktrees: any WorktreeManaging
    private let environment: any TransferImportEnvironment
    private let stagingRoot: URL
    private let now: @Sendable () -> Date

    /// Per-transfer staging state.
    private final class Staging {
        let offer: TransferOffer
        let project: Project
        let sourceDeviceID: String
        let dir: URL
        /// Expected items = the offer's known items plus the always-present bundle.
        var expected: Set<TransferItemKind>
        /// How many chunks of each item have been appended (the resume cursor).
        var received: [TransferItemKind: Int] = [:]
        /// The declared chunk count + whole-item sha, learned from the chunks.
        var chunkCounts: [TransferItemKind: Int] = [:]
        var itemHashes: [TransferItemKind: String] = [:]
        var handles: [TransferItemKind: FileHandle] = [:]
        var complete: Set<TransferItemKind> = []
        var staged = false
        /// The rebuilt destination session, ready to activate on commit.
        var stagedSession: Session?

        init(offer: TransferOffer, project: Project, sourceDeviceID: String, dir: URL) {
            self.offer = offer
            self.project = project
            self.sourceDeviceID = sourceDeviceID
            self.dir = dir
            self.expected = Set(offer.knownItems.map(\.kind)).union([.bundle])
        }
    }
    private var staging: [String: Staging] = [:]

    public init(
        store: any SessionMetadataStore, worktrees: any WorktreeManaging,
        environment: any TransferImportEnvironment, stagingRoot: URL,
        now: @escaping @Sendable () -> Date = { Date() }
    ) {
        self.store = store
        self.worktrees = worktrees
        self.environment = environment
        self.stagingRoot = stagingRoot
        self.now = now
    }

    // MARK: - Offer

    public func receiveOffer(_ offer: TransferOffer, from deviceID: String) async -> [HostMsg] {
        func reject(_ reason: TransferRejectReason, _ message: String? = nil) -> [HostMsg] {
            [.transferReject(TransferReject(transferID: offer.transferID, reason: reason, message: message))]
        }

        // Format + transcript version gates (first consumer of transcript versioning).
        guard offer.formatVersion <= SessionTransferFormat.version else { return reject(.unsupportedFormat) }
        guard offer.transcriptHeaderVersion <= SessionHeader.maxSupported else { return reject(.versionMismatch) }

        // The identifiers below become filesystem paths (staging dir, transcript dir, worktree
        // dir, native-transcript dir) on this Mac, and the offer comes from a paired *peer* — a
        // compromised one could smuggle `../` to write outside our directories. Reject anything
        // that isn't a safe path component before touching the filesystem.
        guard Self.isSafePathComponent(offer.transferID),
              Self.isSafePathComponent(offer.record.sessionID.rawValue),
              Self.isSafeBranch(offer.record.branch),
              // `backendSessionID` names the native-transcript file for memory carry — it too
              // becomes a filesystem path component, so it must be a safe component when present.
              offer.record.backendSessionID.map(Self.isSafePathComponent) ?? true
        else { return reject(.verificationFailed, "Rejected an unsafe identifier in the transfer.") }

        // A re-offer for an already-staging transfer → resume from where we are.
        if let existing = staging[offer.transferID] {
            return [.transferAccept(await accept(for: existing))]
        }

        // Resolve the destination project.
        guard let project = await environment.resolveProject(offer.project) else {
            return reject(.projectNotFound)
        }
        // Branch collision.
        if (try? await worktrees.branchExists(offer.record.branch, in: project)) == true {
            return reject(.branchCollision)
        }
        // One active transfer per session.
        if let _ = try? await store.activeTransfer(sessionID: offer.record.sessionID) {
            return reject(.alreadyTransferring)
        }

        // Claim the inbound lock + open staging.
        let lock = SessionTransferLock(
            transferID: offer.transferID, sessionID: offer.record.sessionID, direction: .inbound,
            peerDeviceID: deviceID, state: .staging, createdAt: now(), updatedAt: now())
        do {
            try await store.beginTransfer(lock)
        } catch {
            return reject(.alreadyTransferring)
        }

        let dir = stagingRoot.appendingPathComponent(offer.transferID, isDirectory: true)
        try? FileManager.default.removeItem(at: dir)
        do {
            try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
        } catch {
            try? await store.clearTransfer(transferID: offer.transferID)
            return reject(.internalError, "Couldn't create staging directory.")
        }
        let entry = Staging(offer: offer, project: project, sourceDeviceID: deviceID, dir: dir)
        staging[offer.transferID] = entry
        return [.transferAccept(await accept(for: entry))]
    }

    private func accept(for entry: Staging) async -> TransferAccept {
        let have = (try? await worktrees.presentSHAs(entry.offer.availableBaseSHAs, in: entry.project)) ?? []
        // Resume cursor: the next chunk index we still need for each item.
        var resume: [TransferItemKind: Int] = [:]
        for kind in entry.expected { resume[kind] = entry.received[kind] ?? 0 }
        return TransferAccept(
            transferID: entry.offer.transferID, resolvedProjectID: entry.project.id,
            haveSHAs: have, resumeFrom: resume)
    }

    // MARK: - Chunk

    public func receiveChunk(_ chunk: TransferChunk) async -> [HostMsg] {
        guard let entry = staging[chunk.transferID] else {
            // A chunk with no live offer — nothing to stage.
            return [.transferReject(TransferReject(
                transferID: chunk.transferID, reason: .internalError, message: "No staged transfer."))]
        }
        entry.chunkCounts[chunk.kind] = chunk.chunkCount
        entry.itemHashes[chunk.kind] = chunk.sha256

        let next = entry.received[chunk.kind] ?? 0
        let ack: HostMsg = .transferChunkAck(TransferChunkAck(
            transferID: chunk.transferID, kind: chunk.kind, chunkIndex: chunk.chunkIndex))
        if chunk.chunkIndex < next {
            return [ack]  // duplicate (resume overlap) — idempotent ack.
        }
        guard chunk.chunkIndex == next else {
            // Out of order — stop-and-wait shouldn't produce this.
            return [.transferReject(TransferReject(
                transferID: chunk.transferID, reason: .verificationFailed, message: "Out-of-order chunk."))]
        }

        do {
            let handle = try openHandle(for: chunk.kind, in: entry)
            try handle.write(contentsOf: chunk.data)
            entry.received[chunk.kind] = next + 1
        } catch {
            return failStaging(entry, .internalError, "Couldn't write staged chunk.")
        }

        // Item complete?
        if entry.received[chunk.kind] == chunk.chunkCount {
            try? entry.handles[chunk.kind]?.close()
            entry.handles[chunk.kind] = nil
            let url = entry.dir.appendingPathComponent(fileName(for: chunk.kind))
            if (try? TransferHash.sha256Hex(ofFileAt: url)) != chunk.sha256 {
                return failStaging(entry, .verificationFailed, "\(chunk.kind.rawValue) hash mismatch.")
            }
            entry.complete.insert(chunk.kind)
        }

        // All items complete → stage the session non-runnable, then report ready.
        if entry.expected.isSubset(of: entry.complete), !entry.staged {
            do {
                try await stage(entry)
                entry.staged = true
                try await store.updateTransferState(transferID: chunk.transferID, to: .ready, at: now())
                return [ack, .transferReady(chunk.transferID)]
            } catch let reject as TransferReject {
                return [ack] + failStaging(entry, reject.reason, reject.message)
            } catch {
                return [ack] + failStaging(entry, .internalError, "Staging failed: \(error)")
            }
        }
        return [ack]
    }

    // MARK: - Commit

    public func receiveCommit(_ transferID: String) async -> [HostMsg] {
        // Idempotent: an already-activated transfer just re-acks.
        if let lock = try? await store.transfer(transferID: transferID), lock.state == .activated {
            return [.transferCommitted(transferID)]
        }
        guard let entry = staging[transferID], entry.staged, let session = entry.stagedSession else {
            return [.transferReject(TransferReject(
                transferID: transferID, reason: .internalError, message: "Nothing staged to activate."))]
        }
        do {
            try await store.activateTransferredSession(session, transferID: transferID)
            await environment.didActivateTransferredSession(session)
            staging[transferID] = nil
            try? FileManager.default.removeItem(at: entry.dir)
            return [.transferCommitted(transferID)]
        } catch {
            return [.transferReject(TransferReject(
                transferID: transferID, reason: .internalError, message: "Activation failed: \(error)"))]
        }
    }

    // MARK: - Cancel

    public func receiveCancel(_ transferID: String) async {
        if let entry = staging[transferID] {
            for handle in entry.handles.values { try? handle.close() }
            try? FileManager.default.removeItem(at: entry.dir)
        }
        staging[transferID] = nil
        try? await store.clearTransfer(transferID: transferID)
    }

    /// Relaunch recovery: drop staging dirs + locks for inbound transfers that never finished
    /// staging (`.staging`). A `.ready` transfer — staged + verified, but the source vanished before
    /// commit — is recovered separately via `recoverableInboundTransfers` (activate-anyway/discard).
    public func discardStaleInboundTransfers() async {
        guard let active = try? await store.allActiveTransfers() else { return }
        for lock in active where lock.direction == .inbound && lock.state == .staging {
            try? FileManager.default.removeItem(
                at: stagingRoot.appendingPathComponent(lock.transferID, isDirectory: true))
            try? await store.clearTransfer(transferID: lock.transferID)
        }
    }

    /// A stranded inbound transfer (mesh P5): staged + verified (`.ready`) on this Mac, but the
    /// source vanished before sending `transferCommit`. Rebuilt from the on-disk manifest after a
    /// relaunch lost the in-memory staging, so the app can offer "Activate anyway" or "Discard".
    public struct RecoverableInboundTransfer: Sendable, Equatable {
        public let transferID: String
        /// The staged destination session (from the manifest) — ready to insert on activate.
        public let session: Session
        /// The source Mac's `HostID` (the inbound lock's peer), for a friendly "from <Mac>" label.
        public let sourceDeviceID: String
    }

    /// Scan for inbound transfers stranded at `.ready` and rebuild each from its persisted manifest.
    /// A `.ready` lock whose manifest/staging is gone is unrecoverable — its orphan lock + dir are
    /// dropped so it doesn't linger. Safe to call repeatedly (launch + on demand).
    public func recoverableInboundTransfers() async -> [RecoverableInboundTransfer] {
        guard let active = try? await store.allActiveTransfers() else { return [] }
        var out: [RecoverableInboundTransfer] = []
        for lock in active where lock.direction == .inbound && lock.state == .ready {
            // Already staged in memory (this session, no relaunch) — recover from that.
            if let entry = staging[lock.transferID], let session = entry.stagedSession {
                out.append(RecoverableInboundTransfer(
                    transferID: lock.transferID, session: session, sourceDeviceID: lock.peerDeviceID))
                continue
            }
            let dir = stagingRoot.appendingPathComponent(lock.transferID, isDirectory: true)
            guard let data = try? Data(contentsOf: dir.appendingPathComponent(Self.stagedManifestName)),
                  let session = try? JSONDecoder().decode(Session.self, from: data) else {
                // Manifest gone → nothing to activate; clear the orphan.
                try? FileManager.default.removeItem(at: dir)
                try? await store.clearTransfer(transferID: lock.transferID)
                continue
            }
            out.append(RecoverableInboundTransfer(
                transferID: lock.transferID, session: session, sourceDeviceID: lock.peerDeviceID))
        }
        return out
    }

    /// Activate a stranded inbound transfer the user chose to keep (mesh P5 "Activate anyway"):
    /// insert the staged session row + flip the lock to `.activated` (atomic, idempotent), notify
    /// the environment (which builds the controller + broadcasts), and clean up the staging dir.
    /// Returns the activated session, or nil if there's nothing recoverable under `transferID`.
    @discardableResult
    public func activateRecoveredTransfer(_ transferID: String) async throws -> Session? {
        let session: Session
        if let entry = staging[transferID], let staged = entry.stagedSession {
            session = staged
        } else {
            let dir = stagingRoot.appendingPathComponent(transferID, isDirectory: true)
            guard let data = try? Data(contentsOf: dir.appendingPathComponent(Self.stagedManifestName)),
                  let decoded = try? JSONDecoder().decode(Session.self, from: data) else { return nil }
            session = decoded
        }
        try await store.activateTransferredSession(session, transferID: transferID)
        await environment.didActivateTransferredSession(session)
        staging[transferID] = nil
        try? FileManager.default.removeItem(
            at: stagingRoot.appendingPathComponent(transferID, isDirectory: true))
        return session
    }

    /// Forget a stranded inbound transfer's in-memory staging + drop its staging dir (mesh P5). The
    /// imported worktree/branch + placed transcript are the app's to tear down (it owns the project +
    /// worktree manager); this just clears what the importer owns, then the lock.
    public func clearInboundStaging(_ transferID: String) async {
        if let entry = staging[transferID] {
            for handle in entry.handles.values { try? handle.close() }
        }
        staging[transferID] = nil
        try? FileManager.default.removeItem(
            at: stagingRoot.appendingPathComponent(transferID, isDirectory: true))
        try? await store.clearTransfer(transferID: transferID)
    }

    // MARK: - Staging (build the non-runnable session)

    private func stage(_ entry: Staging) async throws {
        let record = entry.offer.record
        let project = entry.project
        let slug = record.branch.hasPrefix("nucleic/")
            ? String(record.branch.dropFirst("nucleic/".count)) : record.branch
        let worktreePath = (project.resolvedWorktreeBase as NSString)
            .appendingPathComponent(slug)

        // 1. Import the git branch into a destination worktree.
        let bundleURL = entry.dir.appendingPathComponent(fileName(for: .bundle))
        _ = try await worktrees.createFromTransfer(
            bundle: bundleURL, branch: record.branch, baseSHA: record.baseSHA,
            destPath: worktreePath, in: project)

        // 2. Place our transcript at the destination path and rewrite its header for the new
        //    worktree / native path / memory-carry decision.
        let destTranscriptDir = environment.transcriptsDir
            .appendingPathComponent(record.sessionID.rawValue, isDirectory: true)
        try FileManager.default.createDirectory(at: destTranscriptDir, withIntermediateDirectories: true)
        let destTranscript = destTranscriptDir.appendingPathComponent("transcript.jsonl")
        try? FileManager.default.removeItem(at: destTranscript)
        try FileManager.default.copyItem(
            at: entry.dir.appendingPathComponent(fileName(for: .transcript)), to: destTranscript)

        // 3. Best-effort native memory carry. Only if the source sent it AND the destination can
        //    host it AND the copy succeeds; any failure falls back to fresh context.
        var carriedBackendSessionID: String? = nil
        var destNativePath: String? = nil
        if let backendSessionID = record.backendSessionID,
           entry.complete.contains(.nativeTranscript),
           let destNativeURL = environment.nativeTranscriptDestination(
                session: draftSession(record: record, project: project, worktreePath: worktreePath,
                                      transcriptPath: destTranscript.path, backendSessionID: backendSessionID,
                                      nativeTranscriptPath: nil),
                project: project, worktreePath: worktreePath, backendSessionID: backendSessionID) {
            do {
                try FileManager.default.createDirectory(
                    at: destNativeURL.deletingLastPathComponent(), withIntermediateDirectories: true)
                try? FileManager.default.removeItem(at: destNativeURL)
                try FileManager.default.copyItem(
                    at: entry.dir.appendingPathComponent(fileName(for: .nativeTranscript)), to: destNativeURL)
                carriedBackendSessionID = backendSessionID
                destNativePath = destNativeURL.path
            } catch {
                carriedBackendSessionID = nil  // fresh-context fallback
            }
        }

        // 4. Rewrite the transcript header for the destination.
        let (sourceHeader, _) = try TranscriptReader(url: destTranscript).read()
        let destHeader = TranscriptHeaderRewriter.destinationHeader(
            from: sourceHeader, worktree: worktreePath,
            nativeTranscriptPath: destNativePath, backendSessionID: carriedBackendSessionID)
        try TranscriptHeaderRewriter.rewrite(at: destTranscript, to: destHeader)

        // 5. Build the staged (not-yet-inserted) session. Status is quiescent — a moved session
        //    arrives done/awaiting the user, resumable. Stamp its provenance (where it arrived from)
        //    so the destination can show "Arrived from <Mac>".
        let staged = draftSession(
            record: record, project: project, worktreePath: worktreePath,
            transcriptPath: destTranscript.path, backendSessionID: carriedBackendSessionID,
            nativeTranscriptPath: destNativePath, arrivedFromDeviceID: entry.sourceDeviceID)
        entry.stagedSession = staged

        // 6. Persist the staged session to the staging dir so a destination relaunch (which loses
        //    the in-memory offer + staged session) can still recover it — activate-anyway or discard
        //    — for a transfer stranded at `.ready` because the source vanished before commit. Removed
        //    with the staging dir on a normal commit.
        let manifest = entry.dir.appendingPathComponent(Self.stagedManifestName)
        try JSONEncoder().encode(staged).write(to: manifest, options: .atomic)
    }

    private func draftSession(
        record: SessionTransferRecord, project: Project, worktreePath: String,
        transcriptPath: String, backendSessionID: String?, nativeTranscriptPath: String?,
        arrivedFromDeviceID: String? = nil
    ) -> Session {
        Session(
            id: record.sessionID, projectID: project.id, backend: record.backend,
            backendSessionID: backendSessionID, title: record.title,
            status: record.lastSeq > 0 ? .awaitingInput : .idle,
            worktreePath: worktreePath, branch: record.branch, baseSHA: record.baseSHA,
            model: record.model, effort: record.effort, lastSeq: record.lastSeq,
            transcriptPath: transcriptPath, nativeTranscriptPath: nativeTranscriptPath,
            arrivedFromDeviceID: arrivedFromDeviceID, arrivedAt: arrivedFromDeviceID == nil ? nil : now(),
            createdAt: record.createdAt, updatedAt: now())
    }

    // MARK: - Helpers

    private func openHandle(for kind: TransferItemKind, in entry: Staging) throws -> FileHandle {
        if let handle = entry.handles[kind] { return handle }
        let url = entry.dir.appendingPathComponent(fileName(for: kind))
        FileManager.default.createFile(atPath: url.path, contents: nil)
        let handle = try FileHandle(forWritingTo: url)
        entry.handles[kind] = handle
        return handle
    }

    /// A single path component we'll trust to build a directory/file name: non-empty, no path
    /// separators, no `..`, no leading dot, and only characters git/filesystems produce for our
    /// ids and slugs. Blocks path traversal from a compromised peer's offer.
    static func isSafePathComponent(_ s: String) -> Bool {
        guard !s.isEmpty, s.count <= 255, s != ".", s != ".." else { return false }
        if s.hasPrefix(".") || s.hasPrefix("-") { return false }
        if s.contains("/") || s.contains("\\") || s.contains("\0") || s.contains("..") { return false }
        let allowed = CharacterSet(charactersIn:
            "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_.")
        return s.unicodeScalars.allSatisfy { allowed.contains($0) }
    }

    /// A branch name we'll turn into a worktree dir + a git ref: the `nucleic/<slug>` shape, whose
    /// slug must be a safe component. Rejects anything else (arbitrary refs / traversal).
    static func isSafeBranch(_ branch: String) -> Bool {
        guard branch.hasPrefix("nucleic/") else { return false }
        let slug = String(branch.dropFirst("nucleic/".count))
        return isSafePathComponent(slug)
    }

    private func fileName(for kind: TransferItemKind) -> String {
        switch kind {
        case .bundle: "session.bundle"
        case .transcript: "transcript.jsonl"
        case .nativeTranscript: "native.jsonl"
        }
    }

    /// The staged-session manifest inside a transfer's staging dir — the JSON-encoded destination
    /// `Session`, written at `.ready` so a relaunch can recover a stranded inbound transfer.
    static let stagedManifestName = "staged-session.json"

    private func failStaging(_ entry: Staging, _ reason: TransferRejectReason, _ message: String) -> [HostMsg] {
        for handle in entry.handles.values { try? handle.close() }
        try? FileManager.default.removeItem(at: entry.dir)
        staging[entry.offer.transferID] = nil
        Task { [store, id = entry.offer.transferID] in try? await store.clearTransfer(transferID: id) }
        return [.transferReject(TransferReject(transferID: entry.offer.transferID, reason: reason, message: message))]
    }
}
