import XCTest
import Foundation
@testable import ForgeVideo

final class PNGTests: XCTestCase {

    // PNG signature + chunk layout (IHDR/IDAT/IEND).
    func testStructure() {
        let frame = VideoFrame(
            width: 2, height: 2,
            pixels: [SIMD3(1, 0, 0), SIMD3(0, 1, 0), SIMD3(0, 0, 1), SIMD3(1, 1, 1)])
        let png = PNGWriter.encode(frame)

        // Signature.
        XCTAssertEqual([UInt8](png.prefix(8)), [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A])
        // IHDR right after signature.
        XCTAssertEqual([UInt8](png[12..<16]), [UInt8]("IHDR".utf8))
        // Width/height big-endian in IHDR.
        XCTAssertEqual([UInt8](png[16..<24]), [0, 0, 0, 2, 0, 0, 0, 2])
        // Bit depth 8, color type 2 (truecolor).
        XCTAssertEqual(png[24], 8)
        XCTAssertEqual(png[25], 2)
        // Contains IDAT and ends with IEND.
        XCTAssertNotNil(png.range(of: Data("IDAT".utf8)))
        XCTAssertEqual([UInt8](png.suffix(8).prefix(4)), [UInt8]("IEND".utf8))
    }

    // zlib stream inside IDAT decodes back to original scanlines (via system python zlib).
    func testPixelsRoundTripViaPython() throws {
        let frame = VideoFrame(
            width: 3, height: 2,
            pixels: [
                SIMD3(1, 0, 0), SIMD3(0, 1, 0), SIMD3(0, 0, 1),
                SIMD3(0.5, 0.5, 0.5), SIMD3(0, 0, 0), SIMD3(1, 1, 1),
            ])
        let png = PNGWriter.encode(frame)
        let path = NSTemporaryDirectory() + "forge-png-\(UUID().uuidString).png"
        defer { try? FileManager.default.removeItem(atPath: path) }
        try png.write(to: URL(fileURLWithPath: path))

        // Python: parse IDAT, inflate, check first pixel bytes.
        let script = """
        import struct, zlib, sys
        data = open('\(path)', 'rb').read()
        pos = 8
        idat = b''
        while pos < len(data):
            length, ctype = struct.unpack('>I4s', data[pos:pos+8])
            if ctype == b'IDAT':
                idat += data[pos+8:pos+8+length]
            pos += 12 + length
        raw = zlib.decompress(idat)
        # Row 0: filter byte 0 then RGB bytes.
        assert raw[0] == 0, 'filter'
        assert raw[1:4] == bytes([255, 0, 0]), raw[1:4]
        assert raw[4:7] == bytes([0, 255, 0]), raw[4:7]
        # Row 1 first pixel = 128 gray.
        row1 = raw[1 + 3*3:]
        assert row1[0] == 0, 'filter row1'
        assert row1[1:4] == bytes([128, 128, 128]), row1[1:4]
        print('PNG-OK')
        """
        let proc = Process()
        proc.executableURL = URL(fileURLWithPath: "/usr/bin/python3")
        proc.arguments = ["-c", script]
        let pipe = Pipe()
        proc.standardOutput = pipe
        proc.standardError = pipe
        try proc.run()
        proc.waitUntilExit()
        let out = String(data: pipe.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8) ?? ""
        XCTAssertEqual(proc.terminationStatus, 0, out)
        XCTAssertTrue(out.contains("PNG-OK"), out)
    }

    // Values clamp.
    func testClamping() {
        let frame = VideoFrame(width: 1, height: 1, pixels: [SIMD3(2.0, -1.0, 0.5)])
        let png = PNGWriter.encode(frame)
        XCTAssertGreaterThan(png.count, 40)
    }

    // StillGrabber PNG mode writes .png beside proper name.
    func testStillGrabberPNG() throws {
        let dir = NSTemporaryDirectory() + "forge-pngstill-\(UUID().uuidString)"
        defer { try? FileManager.default.removeItem(atPath: dir) }
        let frame = VideoFrame(width: 4, height: 4, pixels: .init(repeating: SIMD3(0.2, 0.4, 0.8), count: 16))
        let path = try StillGrabber(directory: dir, format: .png).grab(frame: frame, clipName: "R001C001")
        XCTAssertTrue(path.hasSuffix(".png"))
        let data = try Data(contentsOf: URL(fileURLWithPath: path))
        XCTAssertEqual([UInt8](data.prefix(4)), [0x89, 0x50, 0x4E, 0x47])
    }
}