//Licensed to the Apache Software Foundation (ASF) under one
//or more contributor license agreements.  See the NOTICE file
//distributed with this work for additional information
//    regarding copyright ownership.  The ASF licenses this file
//to you under the Apache License, Version 2.0 (the
//"License"); you may not use this file except in compliance
//with the License.  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//Unless required by applicable law or agreed to in writing,
//software distributed under the License is distributed on an
//"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
//KIND, either express or implied.  See the License for the
//specific language governing permissions and limitations
//under the License.

import UIKit
import XCTest
import Network
import AVKit

class UDPConnect {
    
    static let current = UDPConnect()
    
    private var listener: NWListener!
    private var connection: NWConnection!
            
    private let udpThread = DispatchQueue(label: "udpQueue" + UUID().uuidString, qos: .userInitiated)
    
    func stop() {
        listener?.cancel()
        connection?.cancel()
    }
    
    func start() {
        
        // testRTMP()
        
        do {
            let port = UInt16(ATSDevice.current.screenCapturePort)
            listener = try NWListener(using: .udp, on: NWEndpoint.Port(rawValue: port)!)
        } catch {
            print(error.localizedDescription)
            return
        }
        
        listener.newConnectionHandler = { conn in
            self.connection?.cancel()
            self.connection = conn
            
            conn.stateUpdateHandler = { state in
                switch state {
                case .ready:
                    self.receive(on: conn)
                case .failed(let error):
                    print("conn failed : \(error)")
                default:
                    break
                }
            }
            
            conn.start(queue: self.udpThread)
        }
        
        listener.start(queue: DispatchQueue.main)
    }
    
    private func receive(on connection: NWConnection) {
        connection.receiveMessage { (_, _, _, _) in
            
            guard let nextFrame = self.nextFrame() else {
                self.receive(on: connection)
                return
            }

            self.sendFrame(nextFrame, on: connection)
        }
    }
    
    private let packetSize = 2000

    private func sendFrame(_ frame: Data, on connection: NWConnection) {
        var datagramArray: [Data] = []
        var offSet = 0

        repeat {
            let datagramSize = min(packetSize, frame.count - offSet)
            var datagram = frame.subdata(in: offSet..<offSet + datagramSize)
            
            let offsetIndex = UInt32(offSet).toByteArray()
            
            offSet += datagramSize
            let remainingDataCount = UInt32(frame.count - offSet).toByteArray()
            
            datagram.insert(contentsOf: offsetIndex + remainingDataCount, at: 0)
            
            datagramArray.append(datagram)
            
        } while offSet < frame.count
                
        connection.batch {
            datagramArray.forEach { connection.send(content: $0, completion: NWConnection.SendCompletion.contentProcessed { _ in }) }
        }
                
        receive(on: connection)
    }
    
    private func nextFrame() -> Data? {
        let image = XCUIScreen.main.screenshot().image
        
        guard let resizedImage = image.resizeImageTo() else {
            return nil
        }
        
        guard let orientedImage = resizedImage.fixedOrientation() else {
            return nil
        }
        
        guard var data = orientedImage.jpegData(compressionQuality: 0.66) else {
            return nil
        }
        
        let orientation = UInt32(image.imageOrientation.rawValue).toByteArray()
        let offsetX = UInt32(0).toByteArray()
        
        data.insert(contentsOf: orientation + offsetX, at: 0)

        return data
    }
    
    /* private func testRTMP() {
        let rtmpConnection = RTMPConnection()
        let rtmpStream = RTMPStream(connection: rtmpConnection)
        
        rtmpStream.attachScreen(ScreenCaptureSession(shared: UIApplication.shared))

        rtmpConnection.connect("rtmp://localhost/appName/instanceName")
        rtmpStream.publish("streamName")
        // if you want to record a stream.
        // rtmpStream.publish("streamName", type: .localRecord)
    } */
}

extension UInt32 {
    
    func toByteArray() -> [UInt8] {
        var bigEndian = self.bigEndian
        let count = MemoryLayout<Self>.size
        let bytePtr = withUnsafePointer(to: &bigEndian) {
            $0.withMemoryRebound(to: UInt8.self, capacity: count) {
                UnsafeBufferPointer(start: $0, count: count)
            }
        }
        return Array(bytePtr)
    }
}

extension UIImage {
    /// Fix image orientaton to protrait up
    func fixedOrientation() -> UIImage? {
        guard imageOrientation != UIImage.Orientation.up else {
            // This is default orientation, don't need to do anything
            return self.copy() as? UIImage
        }

        guard let cgImage = self.cgImage else {
            // CGImage is not available
            return nil
        }

        guard let colorSpace = cgImage.colorSpace, let ctx = CGContext(data: nil, width: Int(size.width), height: Int(size.height), bitsPerComponent: cgImage.bitsPerComponent, bytesPerRow: 0, space: colorSpace, bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else {
            return nil // Not able to create CGContext
        }

        var transform: CGAffineTransform = CGAffineTransform.identity

        switch imageOrientation {
        case .down, .downMirrored:
            transform = transform.translatedBy(x: size.width, y: size.height)
            transform = transform.rotated(by: CGFloat.pi)
        case .left, .leftMirrored:
            transform = transform.translatedBy(x: size.width, y: 0)
            transform = transform.rotated(by: CGFloat.pi / 2.0)
        case .right, .rightMirrored:
            transform = transform.translatedBy(x: 0, y: size.height)
            transform = transform.rotated(by: CGFloat.pi / -2.0)
        case .up, .upMirrored:
            break
        @unknown default:
            fatalError("Missing...")
            break
        }

        // Flip image one more time if needed to, this is to prevent flipped image
        switch imageOrientation {
        case .upMirrored, .downMirrored:
            transform = transform.translatedBy(x: size.width, y: 0)
            transform = transform.scaledBy(x: -1, y: 1)
        case .leftMirrored, .rightMirrored:
            transform = transform.translatedBy(x: size.height, y: 0)
            transform = transform.scaledBy(x: -1, y: 1)
        case .up, .down, .left, .right:
            break
        @unknown default:
            fatalError("Missing...")
            break
        }

        ctx.concatenate(transform)

        switch imageOrientation {
        case .left, .leftMirrored, .right, .rightMirrored:
            ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: size.height, height: size.width))
        default:
            ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: size.width, height: size.height))
            break
        }

        guard let newCGImage = ctx.makeImage() else { return nil }
        return UIImage.init(cgImage: newCGImage, scale: UIScreen.main.scale, orientation: .up)
    }
}
