import Foundation
import IrdP2PCrypto

/// Splits large application messages into `chunk` inner messages and reassembles
/// them, with bounds on concurrent assemblies, bytes and age.
enum Chunker {
    /// Raw bytes per relayed slice: base64 (×4/3) plus JSON and envelope overhead
    /// (another ×4/3 for the sealed ciphertext) must stay under the frame cap.
    static func relaySliceSize(maxFrameBytes: Int) -> Int {
        min(32_768, max(1024, Int(Double(maxFrameBytes - 512) / 1.85)))
    }

    /// Raw bytes per data-channel slice: keeps each SCTP message far below 16 KiB.
    static let dataChannelSliceSize = 8_192

    /// The chunk id of a payload: the first 8 bytes of
    /// `SHA-256(UTF8(encoding) ‖ UInt32BE(count) ‖ data)`. Derived from the content so a
    /// retransmit of the same payload resumes the receiver's assembly instead of
    /// opening a rival one; identical payloads share an assembly by design.
    static func contentId(_ data: Data, encoding: ChunkEncoding, count: Int) -> String {
        var input = Data(encoding.rawValue.utf8)
        withUnsafeBytes(of: UInt32(count).bigEndian) { input.append(contentsOf: $0) }
        input.append(data)
        return Bytes.base64URL(Bytes.sha256(input).prefix(8))
    }

    static func split(_ data: Data, encoding: ChunkEncoding, sliceSize: Int) -> [InnerMessage] {
        let count = max(1, Int((Double(data.count) / Double(sliceSize)).rounded(.up)))
        let id = contentId(data, encoding: encoding, count: count)
        return (0..<count).map { i in
            let start = data.startIndex + i * sliceSize
            let end = min(start + sliceSize, data.endIndex)
            return .chunk(id: id, index: i, count: count, encoding: encoding, data: Data(data[start..<end]))
        }
    }
}

/// Reassembles one peer's chunks. Violations of the bounds drop the assembly and say why.
struct ChunkAssembler: Sendable {
    static let maxConcurrent = 8

    struct Partial {
        var parts: [Int: Data]
        var count: Int
        var encoding: ChunkEncoding
        var bytes: Int
        var started: Duration
    }

    struct Drop: Equatable {
        var id: String
        var reason: AssemblyDropReason
    }

    let maxBytes: Int
    let maxAge: Duration
    private(set) var partials: [String: Partial] = [:]
    private var totalBytes = 0

    init(maxBytes: Int, maxAge: Duration = .seconds(60)) {
        self.maxBytes = maxBytes
        self.maxAge = maxAge
    }

    enum Result: Equatable { case incomplete, complete(encoding: ChunkEncoding, data: Data), dropped(reason: AssemblyDropReason) }

    /// Adds one slice. Callers `sweep(now:)` first so expired assemblies are reported, not just freed.
    mutating func add(id: String, index: Int, count: Int, encoding: ChunkEncoding, data: Data, now: Duration) -> Result {
        sweep(now: now)
        var p = partials[id] ?? Partial(parts: [:], count: count, encoding: encoding, bytes: 0, started: now)
        guard p.count == count, p.encoding == encoding, index >= 0, index < count else { drop(id); return .dropped(reason: .mismatched) }
        if let have = p.parts[index] {
            // ids are derived from the payload, so a retransmit repeats slices: the same bytes resume the assembly
            if have == data { return .incomplete }
            drop(id)
            return .dropped(reason: .duplicateIndex)
        }
        if partials[id] == nil, partials.count >= ChunkAssembler.maxConcurrent { return .dropped(reason: .tooManyConcurrent) }
        guard totalBytes + data.count <= maxBytes else { drop(id); return .dropped(reason: .tooManyBytes) }
        p.parts[index] = data
        p.bytes += data.count
        totalBytes += data.count
        if p.parts.count == count {
            drop(id)
            var out = Data(capacity: p.bytes)
            for i in 0..<count { out.append(p.parts[i]!) }
            return .complete(encoding: encoding, data: out)
        }
        partials[id] = p
        return .incomplete
    }

    private mutating func drop(_ id: String) {
        if let p = partials.removeValue(forKey: id) { totalBytes -= p.bytes }
    }

    /// Drops assemblies older than `maxAge` and returns them (oldest first).
    @discardableResult
    mutating func sweep(now: Duration) -> [Drop] {
        let expired = partials.filter { now - $0.value.started > maxAge }.sorted { ($0.value.started, $0.key) < ($1.value.started, $1.key) }
        for (id, _) in expired { drop(id) }
        return expired.map { Drop(id: $0.key, reason: .ageExpired) }
    }
}