import Foundation

/// HEIC 文件结构解析和 MakerNotes 二进制注入
/// 用于绕过 CGImageDestination 无法正确写入 Int64 MakerNotes 字段的限制
public enum HEICMakerNoteError: Error, CustomStringConvertible {
    case invalidHEIC(String)
    case metaNotFound
    case iinfNotFound
    case ilocNotFound
    case exifItemNotFound
    case exifLocationNotFound(itemID: UInt32)
    case exifPayloadTooSmall
    case tiffNotFound
    case invalidTIFF(String)
    case exifIFDPointerNotFound
    case makerNoteTagNotFound
    case makerNoteOutOfRange
    case makerNoteTooShort(available: Int, required: Int)

    public var description: String {
        switch self {
        case .invalidHEIC(let msg): return "Invalid HEIC: \(msg)"
        case .metaNotFound: return "meta box not found"
        case .iinfNotFound: return "iinf box not found"
        case .ilocNotFound: return "iloc box not found"
        case .exifItemNotFound: return "Exif item not found"
        case .exifLocationNotFound(let id): return "Exif item location not found for item_ID=\(id)"
        case .exifPayloadTooSmall: return "Exif payload too small"
        case .tiffNotFound: return "TIFF header not found"
        case .invalidTIFF(let msg): return "Invalid TIFF: \(msg)"
        case .exifIFDPointerNotFound: return "ExifIFDPointer (0x8769) not found"
        case .makerNoteTagNotFound: return "MakerNote tag (0x927C) not found"
        case .makerNoteOutOfRange: return "MakerNote data out of range"
        case .makerNoteTooShort(let available, let required):
            return "MakerNote too short: available=\(available), required=\(required)"
        }
    }
}

public final class HEICMakerNotePatcher {

    // MARK: - Public API

    /// 将完整的 MakerNotes 数据注入到 HEIC 文件中
    /// 采用重建 Exif item 的方式，支持扩展 MakerNote 大小
    public static func injectMakerNoteInPlace(fileURL: URL, makerNote: Data) throws {
        var fileData = try Data(contentsOf: fileURL, options: [.mappedIfSafe])

        // 解析文件结构
        let metaRange = try findTopLevelBox(type: "meta", in: fileData)
        guard let metaRange else { throw HEICMakerNoteError.metaNotFound }
        let meta = BoxView(data: fileData, range: metaRange)

        let metaChildrenStart = meta.contentStart + 4
        guard metaChildrenStart <= meta.end else {
            throw HEICMakerNoteError.invalidHEIC("meta content too short")
        }

        guard let iinfRange = try findChildBox(type: "iinf", within: metaChildrenStart..<meta.end, in: fileData) else {
            throw HEICMakerNoteError.iinfNotFound
        }
        guard let ilocRange = try findChildBox(type: "iloc", within: metaChildrenStart..<meta.end, in: fileData) else {
            throw HEICMakerNoteError.ilocNotFound
        }

        let exifItemID = try parseIINFAndFindExifItemID(data: fileData, iinfRange: iinfRange)
        let (locations, ilocInfo) = try parseILOCWithInfo(data: fileData, ilocRange: ilocRange)

        guard let exifLoc = locations[exifItemID] else {
            throw HEICMakerNoteError.exifLocationNotFound(itemID: exifItemID)
        }

        let exifStart = Int(exifLoc.offset)
        let exifLen = Int(exifLoc.length)
        guard exifStart >= 0, exifLen > 0, exifStart + exifLen <= fileData.count else {
            throw HEICMakerNoteError.exifPayloadTooSmall
        }

        // 读取现有 Exif item
        let existingExif = fileData.subdata(in: exifStart..<(exifStart + exifLen))

        // 构建新的 Exif item（替换 MakerNote）
        let newExif = try buildNewExifItem(existingExif: existingExif, newMakerNote: makerNote)

        if newExif.count <= exifLen {
            // 新 Exif 不大于原来的，直接原位替换
            var paddedExif = newExif
            if paddedExif.count < exifLen {
                paddedExif.append(Data(repeating: 0x00, count: exifLen - paddedExif.count))
            }
            fileData.replaceSubrange(exifStart..<(exifStart + exifLen), with: paddedExif)
        } else {
            // 新 Exif 比原来大，需要追加到文件末尾并更新 iloc
            let newExifOffset = fileData.count
            fileData.append(newExif)

            // 更新 iloc 中的 offset 和 length
            try updateILOC(
                in: &fileData,
                ilocRange: ilocRange,
                ilocInfo: ilocInfo,
                itemID: exifItemID,
                newOffset: UInt64(newExifOffset),
                newLength: UInt64(newExif.count)
            )
        }

        try fileData.write(to: fileURL, options: [.atomic])
    }

    // MARK: - Build New Exif Item

    /// 构建新的 Exif item，替换 MakerNote
    private static func buildNewExifItem(existingExif: Data, newMakerNote: Data) throws -> Data {
        guard existingExif.count >= 10 else {
            throw HEICMakerNoteError.exifPayloadTooSmall
        }

        // Exif item 结构：
        // 4 bytes: TIFF header offset (通常是 6，指向 "Exif\0\0" 之后)
        // 4 bytes: "Exif"
        // 2 bytes: \0\0
        // 然后是 TIFF 数据

        let tiffOffsetValue = existingExif.readUInt32BE(at: 0)
        let tiffStart = 4 + Int(tiffOffsetValue)

        guard tiffStart + 8 <= existingExif.count else {
            throw HEICMakerNoteError.tiffNotFound
        }

        // 检查字节序
        let endianMarker = existingExif.subdata(in: tiffStart..<(tiffStart + 2))
        let isBigEndian: Bool
        if endianMarker == Data([0x4D, 0x4D]) {
            isBigEndian = true
        } else if endianMarker == Data([0x49, 0x49]) {
            isBigEndian = false
        } else {
            throw HEICMakerNoteError.invalidTIFF("Invalid endian marker")
        }

        // 构建新的 TIFF 数据（Big-Endian，与 Apple 设备一致）
        var newTiff = Data()

        // TIFF Header: "MM" + 0x002A + IFD0 offset (8)
        newTiff.append(contentsOf: [0x4D, 0x4D])  // Big-endian
        newTiff.append(contentsOf: [0x00, 0x2A])  // TIFF magic
        newTiff.append(contentsOf: [0x00, 0x00, 0x00, 0x08])  // IFD0 offset = 8

        // IFD0: 1 entry (ExifIFDPointer)
        // Entry count: 1
        newTiff.append(contentsOf: [0x00, 0x01])

        // Entry: ExifIFDPointer (0x8769)
        let exifIFDOffset: UInt32 = 8 + 2 + 12 + 4  // = 26 (IFD0 之后)
        newTiff.append(contentsOf: [0x87, 0x69])  // tag
        newTiff.append(contentsOf: [0x00, 0x04])  // type = LONG
        newTiff.append(contentsOf: [0x00, 0x00, 0x00, 0x01])  // count = 1
        newTiff.appendUInt32BE(exifIFDOffset)  // value = offset to Exif IFD

        // Next IFD offset: 0 (no more IFDs)
        newTiff.append(contentsOf: [0x00, 0x00, 0x00, 0x00])

        // Exif IFD: 1 entry (MakerNote)
        let makerNoteDataOffset: UInt32 = exifIFDOffset + 2 + 12 + 4  // = 44
        newTiff.append(contentsOf: [0x00, 0x01])  // entry count

        // Entry: MakerNote (0x927C)
        newTiff.append(contentsOf: [0x92, 0x7C])  // tag
        newTiff.append(contentsOf: [0x00, 0x07])  // type = UNDEFINED
        newTiff.appendUInt32BE(UInt32(newMakerNote.count))  // count
        newTiff.appendUInt32BE(makerNoteDataOffset)  // offset to MakerNote data

        // Next IFD offset: 0
        newTiff.append(contentsOf: [0x00, 0x00, 0x00, 0x00])

        // MakerNote data
        newTiff.append(newMakerNote)

        // 构建完整的 Exif item
        var newExifItem = Data()
        // 4 bytes: offset to TIFF (= 6, 跳过 "Exif\0\0")
        newExifItem.append(contentsOf: [0x00, 0x00, 0x00, 0x06])
        // "Exif\0\0"
        newExifItem.append(contentsOf: [0x45, 0x78, 0x69, 0x66, 0x00, 0x00])
        // TIFF data
        newExifItem.append(newTiff)

        return newExifItem
    }

    // MARK: - Box Parsing

    private struct BoxHeader {
        let type: String
        let size: Int
        let headerSize: Int
        let contentStart: Int
        let end: Int
    }

    private struct BoxView {
        let data: Data
        let range: Range<Int>

        var start: Int { range.lowerBound }
        var end: Int { range.upperBound }

        var header: BoxHeader {
            let size32 = Int(data.readUInt32BE(at: start))
            let type = data.readFourCC(at: start + 4)

            if size32 == 1 {
                let size64 = Int(data.readUInt64BE(at: start + 8))
                return BoxHeader(
                    type: type,
                    size: size64,
                    headerSize: 16,
                    contentStart: start + 16,
                    end: start + size64
                )
            } else if size32 == 0 {
                return BoxHeader(
                    type: type,
                    size: data.count - start,
                    headerSize: 8,
                    contentStart: start + 8,
                    end: data.count
                )
            } else {
                return BoxHeader(
                    type: type,
                    size: size32,
                    headerSize: 8,
                    contentStart: start + 8,
                    end: start + size32
                )
            }
        }

        var contentStart: Int { header.contentStart }
    }

    private static func findTopLevelBox(type: String, in data: Data) throws -> Range<Int>? {
        var cursor = 0
        while cursor + 8 <= data.count {
            let box = try readBoxHeader(at: cursor, data: data)
            if box.type == type { return cursor..<box.end }
            cursor = box.end
        }
        return nil
    }

    private static func findChildBox(type: String, within range: Range<Int>, in data: Data) throws -> Range<Int>? {
        var cursor = range.lowerBound
        while cursor + 8 <= range.upperBound {
            let box = try readBoxHeader(at: cursor, data: data)
            if box.type == type { return cursor..<min(box.end, range.upperBound) }
            cursor = box.end
        }
        return nil
    }

    private static func readBoxHeader(at offset: Int, data: Data) throws -> BoxHeader {
        guard offset + 8 <= data.count else {
            throw HEICMakerNoteError.invalidHEIC("box header out of bounds")
        }
        let size32 = Int(data.readUInt32BE(at: offset))
        let type = data.readFourCC(at: offset + 4)

        if size32 == 1 {
            guard offset + 16 <= data.count else {
                throw HEICMakerNoteError.invalidHEIC("large size box header out of bounds")
            }
            let size64 = Int(data.readUInt64BE(at: offset + 8))
            guard size64 >= 16 else {
                throw HEICMakerNoteError.invalidHEIC("invalid box size")
            }
            return BoxHeader(type: type, size: size64, headerSize: 16, contentStart: offset + 16, end: offset + size64)
        } else if size32 == 0 {
            return BoxHeader(type: type, size: data.count - offset, headerSize: 8, contentStart: offset + 8, end: data.count)
        } else {
            guard size32 >= 8 else {
                throw HEICMakerNoteError.invalidHEIC("invalid box size")
            }
            return BoxHeader(type: type, size: size32, headerSize: 8, contentStart: offset + 8, end: offset + size32)
        }
    }

    // MARK: - iinf / infe Parsing

    private static func parseIINFAndFindExifItemID(data: Data, iinfRange: Range<Int>) throws -> UInt32 {
        let iinf = BoxView(data: data, range: iinfRange).header
        var cursor = iinf.contentStart

        guard cursor + 4 <= iinf.end else {
            throw HEICMakerNoteError.invalidHEIC("iinf too short")
        }
        let version = data.readUInt8(at: cursor)
        cursor += 4

        let entryCount: UInt32
        if version == 0 {
            entryCount = UInt32(data.readUInt16BE(at: cursor))
            cursor += 2
        } else {
            entryCount = data.readUInt32BE(at: cursor)
            cursor += 4
        }

        var foundExif: UInt32?
        var scanned: UInt32 = 0

        while cursor + 8 <= iinf.end, scanned < entryCount {
            let infe = try readBoxHeader(at: cursor, data: data)
            guard infe.type == "infe" else {
                cursor = infe.end
                continue
            }

            var p = infe.contentStart
            guard p + 4 <= infe.end else {
                throw HEICMakerNoteError.invalidHEIC("infe too short")
            }
            let infeVersion = data.readUInt8(at: p)
            p += 4

            let itemID: UInt32
            if infeVersion >= 3 {
                itemID = data.readUInt32BE(at: p); p += 4
            } else {
                itemID = UInt32(data.readUInt16BE(at: p)); p += 2
            }

            p += 2 // item_protection_index

            var itemType = ""
            if infeVersion >= 2 {
                guard p + 4 <= infe.end else {
                    throw HEICMakerNoteError.invalidHEIC("infe item_type out of bounds")
                }
                itemType = data.readFourCC(at: p)
                p += 4
            }

            if itemType == "Exif" {
                foundExif = itemID
                break
            }

            cursor = infe.end
            scanned += 1
        }

        guard let exifID = foundExif else {
            throw HEICMakerNoteError.exifItemNotFound
        }
        return exifID
    }

    // MARK: - iloc Parsing

    private struct ItemLocation {
        let offset: UInt64
        let length: UInt64
    }

    private struct ILOCInfo {
        let version: UInt8
        let offsetSize: Int
        let lengthSize: Int
        let baseOffsetSize: Int
        let indexSize: Int
        let itemEntries: [UInt32: ILOCItemEntry]
    }

    private struct ILOCItemEntry {
        let itemID: UInt32
        let extentOffsetPosition: Int  // 文件中 extent_offset 字段的位置
        let extentLengthPosition: Int  // 文件中 extent_length 字段的位置
    }

    private static func parseILOCWithInfo(data: Data, ilocRange: Range<Int>) throws -> ([UInt32: ItemLocation], ILOCInfo) {
        let iloc = BoxView(data: data, range: ilocRange).header
        var cursor = iloc.contentStart

        guard cursor + 4 <= iloc.end else {
            throw HEICMakerNoteError.invalidHEIC("iloc too short")
        }
        let version = data.readUInt8(at: cursor)
        cursor += 4

        guard cursor + 2 <= iloc.end else {
            throw HEICMakerNoteError.invalidHEIC("iloc header out of bounds")
        }
        let offsetSize = Int(data.readUInt8(at: cursor) >> 4)
        let lengthSize = Int(data.readUInt8(at: cursor) & 0x0F)
        cursor += 1

        let baseOffsetSize = Int(data.readUInt8(at: cursor) >> 4)
        let indexSize = Int(data.readUInt8(at: cursor) & 0x0F)
        cursor += 1

        let itemCount: UInt32
        if version < 2 {
            guard cursor + 2 <= iloc.end else {
                throw HEICMakerNoteError.invalidHEIC("iloc item_count out of bounds")
            }
            itemCount = UInt32(data.readUInt16BE(at: cursor))
            cursor += 2
        } else {
            guard cursor + 4 <= iloc.end else {
                throw HEICMakerNoteError.invalidHEIC("iloc item_count out of bounds")
            }
            itemCount = data.readUInt32BE(at: cursor)
            cursor += 4
        }

        var locations: [UInt32: ItemLocation] = [:]
        var itemEntries: [UInt32: ILOCItemEntry] = [:]

        for _ in 0..<itemCount {
            guard cursor + 2 <= iloc.end else { break }

            let itemID: UInt32
            if version < 2 {
                itemID = UInt32(data.readUInt16BE(at: cursor)); cursor += 2
            } else {
                guard cursor + 4 <= iloc.end else { break }
                itemID = data.readUInt32BE(at: cursor); cursor += 4
            }

            if version == 1 || version == 2 {
                guard cursor + 2 <= iloc.end else { break }
                cursor += 2
            }

            guard cursor + 2 <= iloc.end else { break }
            cursor += 2

            guard cursor + baseOffsetSize <= iloc.end else { break }
            let baseOffset = try data.readUIntBE(at: cursor, size: baseOffsetSize)
            cursor += baseOffsetSize

            guard cursor + 2 <= iloc.end else { break }
            let extentCount = Int(data.readUInt16BE(at: cursor))
            cursor += 2

            var firstExtentOffset: UInt64 = 0
            var firstExtentLength: UInt64 = 0
            var extentOffsetPos = 0
            var extentLengthPos = 0

            for e in 0..<extentCount {
                if (version == 1 || version == 2) && indexSize > 0 {
                    guard cursor + indexSize <= iloc.end else { break }
                    cursor += indexSize
                }

                guard cursor + offsetSize + lengthSize <= iloc.end else { break }

                if e == 0 {
                    extentOffsetPos = cursor
                }
                let extentOffset = try data.readUIntBE(at: cursor, size: offsetSize)
                cursor += offsetSize

                if e == 0 {
                    extentLengthPos = cursor
                }
                let extentLength = try data.readUIntBE(at: cursor, size: lengthSize)
                cursor += lengthSize

                if e == 0 {
                    firstExtentOffset = extentOffset
                    firstExtentLength = extentLength
                }
            }

            let fileOffset = baseOffset + firstExtentOffset
            if firstExtentLength > 0 {
                locations[itemID] = ItemLocation(offset: fileOffset, length: firstExtentLength)
                itemEntries[itemID] = ILOCItemEntry(
                    itemID: itemID,
                    extentOffsetPosition: extentOffsetPos,
                    extentLengthPosition: extentLengthPos
                )
            }
        }

        let info = ILOCInfo(
            version: version,
            offsetSize: offsetSize,
            lengthSize: lengthSize,
            baseOffsetSize: baseOffsetSize,
            indexSize: indexSize,
            itemEntries: itemEntries
        )

        return (locations, info)
    }

    private static func updateILOC(
        in fileData: inout Data,
        ilocRange: Range<Int>,
        ilocInfo: ILOCInfo,
        itemID: UInt32,
        newOffset: UInt64,
        newLength: UInt64
    ) throws {
        guard let entry = ilocInfo.itemEntries[itemID] else {
            throw HEICMakerNoteError.exifLocationNotFound(itemID: itemID)
        }

        // 写入新的 offset
        fileData.writeUIntBE(at: entry.extentOffsetPosition, value: newOffset, size: ilocInfo.offsetSize)

        // 写入新的 length
        fileData.writeUIntBE(at: entry.extentLengthPosition, value: newLength, size: ilocInfo.lengthSize)
    }

    // MARK: - EXIF/TIFF Patching

    enum Endian {
        case little
        case big
    }
}

// MARK: - Data Extensions

private extension Data {
    func readUInt8(at offset: Int) -> UInt8 {
        self[self.index(self.startIndex, offsetBy: offset)]
    }

    func readUInt16BE(at offset: Int) -> UInt16 {
        let b0 = UInt16(readUInt8(at: offset))
        let b1 = UInt16(readUInt8(at: offset + 1))
        return (b0 << 8) | b1
    }

    func readUInt32BE(at offset: Int) -> UInt32 {
        let b0 = UInt32(readUInt8(at: offset))
        let b1 = UInt32(readUInt8(at: offset + 1))
        let b2 = UInt32(readUInt8(at: offset + 2))
        let b3 = UInt32(readUInt8(at: offset + 3))
        return (b0 << 24) | (b1 << 16) | (b2 << 8) | b3
    }

    func readUInt64BE(at offset: Int) -> UInt64 {
        var v: UInt64 = 0
        for i in 0..<8 {
            v = (v << 8) | UInt64(readUInt8(at: offset + i))
        }
        return v
    }

    func readFourCC(at offset: Int) -> String {
        let bytes = self.subdata(in: offset..<(offset + 4))
        return String(bytes: bytes, encoding: .ascii) ?? "????"
    }

    func readUIntBE(at offset: Int, size: Int) throws -> UInt64 {
        if size == 0 { return 0 }
        guard offset + size <= count else {
            throw HEICMakerNoteError.invalidHEIC("Variable-length integer out of bounds")
        }
        var v: UInt64 = 0
        for i in 0..<size {
            v = (v << 8) | UInt64(readUInt8(at: offset + i))
        }
        return v
    }

    mutating func appendUInt32BE(_ value: UInt32) {
        append(UInt8((value >> 24) & 0xFF))
        append(UInt8((value >> 16) & 0xFF))
        append(UInt8((value >> 8) & 0xFF))
        append(UInt8(value & 0xFF))
    }

    mutating func writeUIntBE(at offset: Int, value: UInt64, size: Int) {
        for i in 0..<size {
            let byteIndex = self.index(self.startIndex, offsetBy: offset + i)
            let shift = (size - 1 - i) * 8
            self[byteIndex] = UInt8((value >> shift) & 0xFF)
        }
    }
}