import Foundation
import AVFoundation

class MP4Box {
    static func create(_ data: Data) throws -> MP4Box {
        let buffer: ByteArray = ByteArray(data: data)
        let size: UInt32 = try buffer.readUInt32()
        let type: String = try buffer.readUTF8Bytes(4)

        buffer.clear()
        switch type {
        case "moov", "trak", "mdia", "minf", "stbl", "edts":
            return MP4ContainerBox(size: size, type: type)
        case "mp4v", "s263", "avc1":
            return MP4VisualSampleEntryBox(size: size, type: type)
        case "mvhd", "mdhd":
            return MP4MediaHeaderBox(size: size, type: type)
        case "mp4a":
            return MP4AudioSampleEntryBox(size: size, type: type)
        case "esds":
            return MP4ElementaryStreamDescriptorBox(size: size, type: type)
        case "stts":
            return MP4TimeToSampleBox(size: size, type: type)
        case "stss":
            return MP4SyncSampleBox(size: size, type: type)
        case "stsd":
            return MP4SampleDescriptionBox(size: size, type: type)
        case "stco":
            return MP4ChunkOffsetBox(size: size, type: type)
        case "stsc":
            return MP4SampleToChunkBox(size: size, type: type)
        case "stsz":
            return MP4SampleSizeBox(size: size, type: type)
        case "elst":
            return MP4EditListBox(size: size, type: type)
        default:
            return MP4Box(size: size, type: type)
        }
    }

    var leafNode: Bool {
        return false
    }

    fileprivate(set) var type: String = "undf"
    fileprivate(set) var size: UInt32 = 0
    fileprivate(set) var offset: UInt64 = 0
    fileprivate(set) var parent: MP4Box?

    init() {
    }

    init(size: UInt32, type: String) {
        self.size = size
        self.type = type
    }

    func load(_ fileHandle: FileHandle) throws -> UInt32 {
        if size == 0 {
            size = UInt32(fileHandle.seekToEndOfFile() - offset)
            return size
        }
        fileHandle.seek(toFileOffset: offset + UInt64(size))
        return size
    }

    func getBoxes(byName: String) -> [MP4Box] {
        return []
    }

    func clear() {
        parent = nil
    }

    func create(_ data: Data, offset: UInt32) throws -> MP4Box {
        let box: MP4Box = try MP4Box.create(data)
        box.parent = self
        box.offset = self.offset + UInt64(offset)
        return box
    }
}

extension MP4Box: CustomStringConvertible {
    // MARK: CustomStringConvertible
    var description: String {
        return Mirror(reflecting: self).description
    }
}

// MARK: -
class MP4ContainerBox: MP4Box {

    fileprivate var children: [MP4Box] = []

    override var leafNode: Bool {
        return false
    }

    override func load(_ file: FileHandle) throws -> UInt32 {
        children.removeAll(keepingCapacity: false)

        var offset: UInt32 = parent == nil ? 0 : 8
        file.seek(toFileOffset: self.offset + UInt64(offset))

        while size != offset {
            let child: MP4Box = try create(file.readData(ofLength: 8), offset: offset)
            offset += try child.load(file)
            children.append(child)
        }

        return offset
    }

    override func getBoxes(byName: String) -> [MP4Box] {
        var list: [MP4Box] = []
        for child in children {
            if byName == child.type || byName == "*" {
                list.append(child)
            }
            if !child.leafNode {
                list += child.getBoxes(byName: byName)
            }
        }
        return list
    }

    override func clear() {
        for child in children {
            child.clear()
        }
        children.removeAll(keepingCapacity: false)
        parent = nil
    }
}

// MARK: -
final class MP4MediaHeaderBox: MP4Box {
    var version: UInt8 = 0
    var creationTime: UInt32 = 0
    var modificationTime: UInt32 = 0
    var timeScale: UInt32 = 0
    var duration: UInt32 = 0
    var language: UInt16 = 0
    var quality: UInt16 = 0

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(size) - 8))
        version = try buffer.readUInt8()
        buffer.position += 3
        creationTime = try buffer.readUInt32()
        modificationTime = try buffer.readUInt32()
        timeScale = try buffer.readUInt32()
        duration = try buffer.readUInt32()
        language = try buffer.readUInt16()
        quality = try buffer.readUInt16()
        buffer.clear()
        return size
    }
}

// MARK: -
final class MP4ChunkOffsetBox: MP4Box {
    var entries: [UInt32] = []

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(size) - 8))
        buffer.position += 4

        let numberOfEntries: UInt32 = try buffer.readUInt32()
        for _ in 0..<numberOfEntries {
            entries.append(try buffer.readUInt32())
        }
        buffer.clear()

        return size
    }
}

// MARK: -
final class MP4SyncSampleBox: MP4Box {
    var entries: [UInt32] = []

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        entries.removeAll(keepingCapacity: false)

        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(size) - 8))
        buffer.position += 4

        let numberOfEntries: UInt32 = try buffer.readUInt32()
        for _ in 0..<numberOfEntries {
            entries.append(try buffer.readUInt32())
        }

        return size
    }
}

// MARK: -
final class MP4TimeToSampleBox: MP4Box {
    struct Entry: CustomStringConvertible {
        var sampleCount: UInt32 = 0
        var sampleDuration: UInt32 = 0

        var description: String {
            return Mirror(reflecting: self).description
        }

        init(sampleCount: UInt32, sampleDuration: UInt32) {
            self.sampleCount = sampleCount
            self.sampleDuration = sampleDuration
        }
    }

    var entries: [Entry] = []

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        entries.removeAll(keepingCapacity: false)

        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(size) - 8))
        buffer.position += 4

        let numberOfEntries: UInt32 = try buffer.readUInt32()
        for _ in 0..<numberOfEntries {
            entries.append(Entry(
                sampleCount: try buffer.readUInt32(),
                sampleDuration: try buffer.readUInt32()
            ))
        }

        return size
    }
}

// MARK: -
final class MP4SampleSizeBox: MP4Box {
    var entries: [UInt32] = []

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        entries.removeAll(keepingCapacity: false)

        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(self.size) - 8))
        buffer.position += 4

        let sampleSize: UInt32 = try buffer.readUInt32()
        if sampleSize != 0 {
            entries.append(sampleSize)
            return size
        }

        let numberOfEntries: UInt32 = try buffer.readUInt32()
        for _ in 0..<numberOfEntries {
            entries.append(try buffer.readUInt32())
        }
        buffer.clear()

        return size
    }
}

// MARK: -
final class MP4ElementaryStreamDescriptorBox: MP4ContainerBox {
    var audioDecorderSpecificConfig: Data = Data()

    var tag: UInt8 = 0
    var tagSize: UInt8 = 0
    var id: UInt16 = 0
    var streamDependenceFlag: UInt8 = 0
    var urlFlag: UInt8 = 0
    var ocrStreamFlag: UInt8 = 0
    var streamPriority: UInt8 = 0

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        var tagSize: UInt8 = 0
        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(self.size) - 8))
        buffer.position += 4

        tag = try buffer.readUInt8()
        self.tagSize = try buffer.readUInt8()
        if self.tagSize == 0x80 {
            buffer.position += 2
            self.tagSize = try buffer.readUInt8()
        }

        id = try buffer.readUInt16()

        let data: UInt8 = try buffer.readUInt8()
        streamDependenceFlag = data >> 7
        urlFlag = (data >> 6) & 0x1
        ocrStreamFlag = (data >> 5) & 0x1
        streamPriority = data & 0x1f

        if streamDependenceFlag == 1 {
            let _: UInt16 = try buffer.readUInt16()
        }

        // Decorder Config Descriptor
        let _: UInt8 = try buffer.readUInt8()
        tagSize = try buffer.readUInt8()
        if tagSize == 0x80 {
            buffer.position += 2
            tagSize = try buffer.readUInt8()
        }
        buffer.position += 13

        // Audio Decorder Spec Info
        let _: UInt8 = try buffer.readUInt8()
        tagSize = try buffer.readUInt8()
        if tagSize == 0x80 {
            buffer.position += 2
            tagSize = try buffer.readUInt8()
        }

        audioDecorderSpecificConfig = try buffer.readBytes(Int(tagSize))

        return size
    }
}

// MARK: -
final class MP4AudioSampleEntryBox: MP4ContainerBox {
    var version: UInt16 = 0

    var channelCount: UInt16 = 0
    var sampleSize: UInt16 = 0
    var compressionId: UInt16 = 0
    var packetSize: UInt16 = 0
    var sampleRate: UInt32 = 0
    var samplesPerPacket: UInt32 = 0
    var bytesPerPacket: UInt32 = 0
    var bytesPerFrame: UInt32 = 0
    var bytesPerSample: UInt32 = 0

    var soundVersion2Data: [UInt8] = []

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(size) - 8))
        buffer.position += 8

        version = try buffer.readUInt16()
        buffer.position += 6

        channelCount = try buffer.readUInt16()
        sampleSize = try buffer.readUInt16()
        compressionId = try buffer.readUInt16()
        packetSize = try buffer.readUInt16()
        sampleRate = try buffer.readUInt32()

        if type != "mlpa" {
            sampleRate = sampleRate >> 16
        }

        if 0 < version {
            samplesPerPacket = try buffer.readUInt32()
            bytesPerPacket = try buffer.readUInt32()
            bytesPerFrame = try buffer.readUInt32()
            bytesPerSample = try buffer.readUInt32()
        }

        if version == 2 {
            soundVersion2Data += try buffer.readBytes(20)
        }

        var offset: UInt32 = UInt32(buffer.position) + 8
        fileHandle.seek(toFileOffset: self.offset + UInt64(offset))

        let esds: MP4Box = try create(fileHandle.readData(ofLength: 8), offset: offset)
        offset += try esds.load(fileHandle)
        children.append(esds)

        // skip
        fileHandle.seek(toFileOffset: self.offset + UInt64(size))

        return size
    }
}

// MARK: -
final class MP4VisualSampleEntryBox: MP4ContainerBox {
    static var dataSize: Int = 78

    var width: UInt16 = 0
    var height: UInt16 = 0
    var hSolution: UInt32 = 0
    var vSolution: UInt32 = 0
    var frameCount: UInt16 = 1
    var compressorname: String = ""
    var depth: UInt16 = 16

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: MP4VisualSampleEntryBox.dataSize))

        buffer.position += 24
        width = try buffer.readUInt16()
        height = try buffer.readUInt16()
        hSolution = try buffer.readUInt32()
        vSolution = try buffer.readUInt32()
        buffer.position += 4
        frameCount = try buffer.readUInt16()
        compressorname = try buffer.readUTF8Bytes(32)
        depth = try buffer.readUInt16()
        let _: UInt16 = try buffer.readUInt16()
        buffer.clear()

        var offset: UInt32 = UInt32(MP4VisualSampleEntryBox.dataSize)
        let child: MP4Box = try MP4Box.create(fileHandle.readData(ofLength: 8))
        child.parent = self
        child.offset = self.offset + UInt64(offset) + 8
        offset += try child.load(fileHandle)
        children.append(child)

        // skip
        fileHandle.seek(toFileOffset: self.offset + UInt64(size))

        return size
    }
}

// MARK: -
final class MP4SampleDescriptionBox: MP4ContainerBox {
    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        children.removeAll(keepingCapacity: false)

        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: 8))
        buffer.position = 4

        var offset: UInt32 = 16
        let numberOfEntries: UInt32 = try buffer.readUInt32()
        for _ in 0..<numberOfEntries {
            let child: MP4Box = try create(fileHandle.readData(ofLength: 8), offset: offset)
            offset += try child.load(fileHandle)
            children.append(child)
        }

        return offset
    }
}

// MARK: -
final class MP4SampleToChunkBox: MP4Box {
    struct Entry: CustomStringConvertible {
        var firstChunk: UInt32 = 0
        var samplesPerChunk: UInt32 = 0
        var sampleDescriptionIndex: UInt32 = 0

        var description: String {
            return Mirror(reflecting: self).description
        }

        init(firstChunk: UInt32, samplesPerChunk: UInt32, sampleDescriptionIndex: UInt32) {
            self.firstChunk = firstChunk
            self.samplesPerChunk = samplesPerChunk
            self.sampleDescriptionIndex = sampleDescriptionIndex
        }
    }

    var entries: [Entry] = []

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(size) - 8))
        buffer.position += 4

        let numberOfEntries: UInt32 = try buffer.readUInt32()
        for _ in 0..<numberOfEntries {
            entries.append(Entry(
                firstChunk: try buffer.readUInt32(),
                samplesPerChunk: try buffer.readUInt32(),
                sampleDescriptionIndex: try buffer.readUInt32()
            ))
        }
        buffer.clear()

        return size
    }
}

// MARK: -
final class MP4EditListBox: MP4Box {
    struct Entry: CustomStringConvertible {
        var segmentDuration: UInt32 = 0
        var mediaTime: UInt32 = 0
        var mediaRate: UInt32 = 0

        var description: String {
            return Mirror(reflecting: self).description
        }

        init(segmentDuration: UInt32, mediaTime: UInt32, mediaRate: UInt32) {
            self.segmentDuration = segmentDuration
            self.mediaTime = mediaTime
            self.mediaRate = mediaRate
        }
    }

    var version: UInt32 = 0
    var entries: [Entry] = []

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        let buffer: ByteArray = ByteArray(data: fileHandle.readData(ofLength: Int(size) - 8))

        version = try buffer.readUInt32()
        entries.removeAll(keepingCapacity: false)

        let numberOfEntries: UInt32 = try buffer.readUInt32()
        for _ in 0..<numberOfEntries {
            entries.append(Entry(
                segmentDuration: try buffer.readUInt32(),
                mediaTime: try buffer.readUInt32(),
                mediaRate: try buffer.readUInt32()
            ))
        }

        return size
    }
}

// MARK: -
final class MP4Reader: MP4ContainerBox {
    private(set) var url: URL

    var isEmpty: Bool {
        return getBoxes(byName: "mdhd").isEmpty
    }

    private var fileHandle: FileHandle?

    init(url: URL) {
        do {
            self.url = url
            super.init()
            fileHandle = try FileHandle(forReadingFrom: url)
        } catch let error as NSError {
            print("\(error)")
        }
    }

    func seek(toFileOffset: UInt64) {
        return fileHandle!.seek(toFileOffset: toFileOffset)
    }

    func readData(ofLength: Int) -> Data {
        return fileHandle!.readData(ofLength: ofLength)
    }

    func readData(ofBox: MP4Box) -> Data {
        guard let fileHandle: FileHandle = fileHandle else {
            return Data()
        }
        let currentOffsetInFile: UInt64 = fileHandle.offsetInFile
        fileHandle.seek(toFileOffset: ofBox.offset + 8)
        let data: Data = fileHandle.readData(ofLength: Int(ofBox.size) - 8)
        fileHandle.seek(toFileOffset: currentOffsetInFile)
        return data
    }

    func load() throws -> UInt32 {
        guard let fileHandle: FileHandle = self.fileHandle else {
            return 0
        }
        return try load(fileHandle)
    }

    func close() {
        fileHandle?.closeFile()
    }

    override func load(_ fileHandle: FileHandle) throws -> UInt32 {
        let size: UInt64 = fileHandle.seekToEndOfFile()
        fileHandle.seek(toFileOffset: 0)
        self.size = UInt32(size)
        return try super.load(fileHandle)
    }
}

// MARK: -
final class MP4TrakReader {
    static let defaultBufferTime: Double = 500

    var trak: MP4Box
    var bufferTime: Double = MP4TrakReader.defaultBufferTime
    weak var delegate: MP4SamplerDelegate?

    private var id: Int = 0
    private var handle: FileHandle?
    private lazy var timerDriver: TimerDriver = {
        var timerDriver: TimerDriver = TimerDriver()
        timerDriver.delegate = self
        return timerDriver
    }()
    private var currentOffset: UInt64 {
        return UInt64(offset[cursor])
    }
    private var currentIsKeyframe: Bool {
        return keyframe[cursor] != nil
    }
    private var currentDuration: Double {
        return Double(totalTimeToSample) * 1000 / Double(timeScale)
    }
    private var currentTimeToSample: Double {
        return Double(timeToSample[cursor]) * 1000 / Double(timeScale)
    }
    private var currentSampleSize: Int {
        return Int((sampleSize.count == 1) ? sampleSize[0] : sampleSize[cursor])
    }
    private var cursor: Int = 0
    private var offset: [UInt32] = []
    private var keyframe: [Int: Bool] = [: ]
    private var timeScale: UInt32 = 0
    private var sampleSize: [UInt32] = []
    private var timeToSample: [UInt32] = []
    private var totalTimeToSample: UInt32 = 0

    init(id: Int, trak: MP4Box) {
        self.id = id
        self.trak = trak

        let mdhd: MP4Box? = trak.getBoxes(byName: "mdhd").first
        if let mdhd: MP4MediaHeaderBox = mdhd as? MP4MediaHeaderBox {
            timeScale = mdhd.timeScale
        }

        let stss: MP4Box? = trak.getBoxes(byName: "stss").first
        if let stss: MP4SyncSampleBox = stss as? MP4SyncSampleBox {
            var keyframes: [UInt32] = stss.entries
            for i in 0..<keyframes.count {
                keyframe[Int(keyframes[i]) - 1] = true
            }
        }

        let stts: MP4Box? = trak.getBoxes(byName: "stts").first
        if let stts: MP4TimeToSampleBox = stts as? MP4TimeToSampleBox {
            var timeToSample: [MP4TimeToSampleBox.Entry] = stts.entries
            for i in 0..<timeToSample.count {
                let entry: MP4TimeToSampleBox.Entry = timeToSample[i]
                for _ in 0..<entry.sampleCount {
                    self.timeToSample.append(entry.sampleDuration)
                }
            }
        }

        let stsz: MP4Box? = trak.getBoxes(byName: "stsz").first
        if let stsz: MP4SampleSizeBox = stsz as? MP4SampleSizeBox {
            sampleSize = stsz.entries
        }

        let stco: MP4Box = trak.getBoxes(byName: "stco").first!
        let stsc: MP4Box = trak.getBoxes(byName: "stsc").first!
        var offsets: [UInt32] = (stco as! MP4ChunkOffsetBox).entries
        var sampleToChunk: [MP4SampleToChunkBox.Entry] = (stsc as! MP4SampleToChunkBox).entries

        var index: Int = 0
        let count: Int = sampleToChunk.count
        for i in 0..<count {
            let m: Int = (i + 1 < count) ? Int(sampleToChunk[i + 1].firstChunk) - 1 : offsets.count
            for j in (Int(sampleToChunk[i].firstChunk) - 1)..<m {
                var offset: UInt32 = offsets[j]
                for _ in 0..<sampleToChunk[i].samplesPerChunk {
                    self.offset.append(offset)
                    offset += sampleSize[index]
                    index += 1
                }
            }
        }
        totalTimeToSample = timeToSample[cursor]
    }

    func execute(_ reader: MP4Reader) {
        do {
            handle = try FileHandle(forReadingFrom: reader.url)

            if let avcC: MP4Box = trak.getBoxes(byName: "avcC").first {
                delegate?.didSet(config: reader.readData(ofBox: avcC), withID: id, type: .video)
            }
            if let esds: MP4ElementaryStreamDescriptorBox = trak.getBoxes(byName: "esds").first as? MP4ElementaryStreamDescriptorBox {
                delegate?.didSet(config: Data(esds.audioDecorderSpecificConfig), withID: id, type: .audio)
            }

            timerDriver.interval = MachUtil.nanosToAbs(UInt64(currentTimeToSample * 1000 * 1000))
            while currentDuration <= bufferTime {
                tick(timerDriver)
            }
            timerDriver.startRunning()
        } catch {
            print("file open error : \(reader.url)")
        }
    }

    private func hasNext() -> Bool {
        return cursor + 1 < offset.count
    }

    private func next() {
        defer {
            cursor += 1
        }
        totalTimeToSample += timeToSample[cursor]
    }
}

extension MP4TrakReader: TimerDriverDelegate {
    // MARK: TimerDriverDelegate
    func tick(_ driver: TimerDriver) {
        guard let handle: FileHandle = handle else {
            driver.stopRunning()
            return
        }
        driver.interval = MachUtil.nanosToAbs(UInt64(currentTimeToSample * 1000 * 1000))
        handle.seek(toFileOffset: currentOffset)
        autoreleasepool {
            delegate?.output(
                data: handle.readData(ofLength: currentSampleSize),
                withID: id,
                currentTime: currentTimeToSample,
                keyframe: currentIsKeyframe
            )
        }
        if hasNext() {
            next()
        } else {
            driver.stopRunning()
        }
    }
}

extension MP4TrakReader: CustomStringConvertible {
    // MARK: CustomStringConvertible
    var description: String {
        return Mirror(reflecting: self).description
    }
}