import HTTP
import Vapor
import Foundation
import ZIPFoundation

/// Server Gzip middleware:
/// 1. checks if the "Accept-Encoding" header contains "gzip"
/// 2. if so, compresses the body and sets the response header "Content-Encoding" to "gzip",
public struct GzipServerMiddleware: Middleware, ServiceType {
    public static func makeService(for worker: Container) throws -> GzipServerMiddleware {
        return .init()
    }

    private let shouldGzipRequest: (_ request: Request) -> Bool
    private let shouldGzipResponse: (_ response: Response) -> Bool

    /// The `shouldGzip` closure is asked for every request whether that request
    /// should allow response gzipping. Returns `true` always by default.
    public init(shouldGzip: @escaping (_ request: Request) -> Bool = { _ in true }) {
        self.shouldGzipRequest = shouldGzip
        self.shouldGzipResponse = { _ in true }
    }

    /// The `shouldGzipRequest` closure is asked for every request whether that request
    /// should allow response gzipping.  `shouldGzipResponse` asks the same for the response.
    /// Both return`true` always by default.
    public init(shouldGzipRequest: @escaping (_ request: Request) -> Bool = { _ in true },
                shouldGzipResponse: @escaping (_ response: Response) -> Bool = { _ in true }
    ) {
        self.shouldGzipRequest = shouldGzipRequest
        self.shouldGzipResponse = shouldGzipResponse
    }

    public func respond(to request: Request, chainingTo next: Responder) throws -> EventLoopFuture<Response> {
        let acceptsGzip = request.http.headers[.acceptEncoding].first?.contains("gzip") ?? false
        let response = try next.respond(to: request)
        guard acceptsGzip && shouldGzipRequest(request) else {
            return response
        }

        return response.flatMap { response in
            guard self.shouldGzipResponse(response) else {
                return request.future(response)
            }
            var headers = response.http.headers
            headers.replaceOrAdd(name: .contentEncoding, value: "gzip")
            return response.http.body.consumeData(on: request).map { data in
                let stream  = HTTPChunkedStream(on: request)
                let bufSize = 4096
                var buffer  = ByteBufferAllocator().buffer(capacity: 16)
                let header  : [UInt8] = [0x1f, 0x8b, 0x08, 0x00]
                let header2 : [UInt8] = [0x00, 0x03]
                buffer.write(bytes: header)
                buffer.write(integer: UInt32(Date().timeIntervalSince1970), endianness: .little)
                buffer.write(bytes: header2)
                var write   = stream.write(.chunk(buffer))
                let crc32   = try Data.compress(size: data.count, bufferSize: bufSize,
                    provider: {(offset, readSize) -> Data in
                        return data.subdata(in: offset..<offset+readSize)
                    },
                    consumer: { data -> Void in
		    	guard data.count > 0 else { return } // Skip empty buffers
                        var buffer  = ByteBufferAllocator().buffer(capacity: bufSize)
                        buffer.write(bytes: data)
                        write = write.flatMap {
                            return stream.write(.chunk(buffer))
                        }
                    })
                write = write.flatMap {
                    buffer.clear()
                    buffer.write(integer: crc32, endianness: .little)
                    buffer.write(integer: UInt32(data.count), endianness: .little)
                    return stream.write(.chunk(buffer))
                }.flatMap {
                    return stream.write(.end)
                }
                DispatchQueue.global().async {
                    _ = try? write.wait()
                }

                let httpResponse = HTTPResponse(status: response.http.status, headers: headers, body: stream)
                return request.response(http: httpResponse)
            }
        }
    }
}