import Foundation

/// Minimal dependency-free PNG encoder: 8-bit truecolor, stored (uncompressed)
/// deflate blocks inside a valid zlib stream. Bigger files than real deflate,
/// but universally decodable — browsers, ffmpeg, Python all accept it.
/// Purpose: shot-report thumbnails + still grabs without an image library.
public enum PNGWriter {

    // MARK: CRC32 (PNG chunk checksums)

    private static let crcTable: [UInt32] = {
        (0..<256).map { n -> UInt32 in
            var c = UInt32(n)
            for _ in 0..<8 {
                c = (c & 1) != 0 ? 0xEDB88320 ^ (c >> 1) : c >> 1
            }
            return c
        }
    }()

    static func crc32(_ data: Data) -> UInt32 {
        var c: UInt32 = 0xFFFFFFFF
        for byte in data {
            c = crcTable[Int((c ^ UInt32(byte)) & 0xFF)] ^ (c >> 8)
        }
        return c ^ 0xFFFFFFFF
    }

    // MARK: Adler32 (zlib checksum)

    static func adler32(_ data: Data) -> UInt32 {
        var a: UInt32 = 1, b: UInt32 = 0
        for byte in data {
            a = (a + UInt32(byte)) % 65521
            b = (b + a) % 65521
        }
        return (b << 16) | a
    }

    private static func putU32BE(_ v: UInt32, into data: inout Data) {
        data.append(UInt8((v >> 24) & 0xFF))
        data.append(UInt8((v >> 16) & 0xFF))
        data.append(UInt8((v >> 8) & 0xFF))
        data.append(UInt8(v & 0xFF))
    }

    private static func chunk(_ type: String, _ payload: Data) -> Data {
        var out = Data()
        putU32BE(UInt32(payload.count), into: &out)
        var typed = Data(type.utf8)
        typed.append(payload)
        out.append(typed)
        putU32BE(crc32(typed), into: &out)
        return out
    }

    /// zlib stream with stored deflate blocks (BTYPE=00).
    static func zlibStored(_ raw: Data) -> Data {
        var out = Data([0x78, 0x01]) // CMF/FLG: 32K window, no preset, check bits valid
        var offset = 0
        let maxBlock = 65535
        while offset < raw.count {
            let size = min(maxBlock, raw.count - offset)
            let isLast: UInt8 = (offset + size >= raw.count) ? 1 : 0
            out.append(isLast) // BFINAL + BTYPE=00
            out.append(UInt8(size & 0xFF))
            out.append(UInt8((size >> 8) & 0xFF))
            out.append(UInt8(~size & 0xFF))
            out.append(UInt8((~size >> 8) & 0xFF))
            out.append(raw.subdata(in: raw.startIndex + offset..<raw.startIndex + offset + size))
            offset += size
        }
        putU32BE(adler32(raw), into: &out)
        return out
    }

    /// Encode frame as 8-bit RGB PNG.
    public static func encode(_ frame: VideoFrame) -> Data {
        var png = Data([0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A])

        // IHDR
        var ihdr = Data()
        putU32BE(UInt32(frame.width), into: &ihdr)
        putU32BE(UInt32(frame.height), into: &ihdr)
        ihdr.append(contentsOf: [8, 2, 0, 0, 0]) // depth 8, truecolor, deflate, adaptive, no interlace
        png.append(chunk("IHDR", ihdr))

        // Scanlines: filter byte 0 + RGB triples.
        var raw = Data(capacity: frame.height * (1 + frame.width * 3))
        for y in 0..<frame.height {
            raw.append(0)
            for x in 0..<frame.width {
                let px = frame.pixels[y * frame.width + x]
                raw.append(UInt8(min(max(px.x, 0), 1) * 255 + 0.5))
                raw.append(UInt8(min(max(px.y, 0), 1) * 255 + 0.5))
                raw.append(UInt8(min(max(px.z, 0), 1) * 255 + 0.5))
            }
        }
        png.append(chunk("IDAT", zlibStored(raw)))
        png.append(chunk("IEND", Data()))
        return png
    }
}