import ChibiRay
import JavaScriptKit
import JavaScriptEventLoop

JavaScriptEventLoop.installGlobalExecutor()
WebWorkerTaskExecutor.installGlobalExecutor()

func renderInCanvas(ctx: JSObject, image: ImageView) {
    let imageData = ctx.createImageData!(image.width, image.height).object!
    let data = imageData.data.object!
    
    for y in 0..<image.height {
        for x in 0..<image.width {
            let index = (y * image.width + x) * 4
            let pixel = image[x, y]
            data[index]     = .number(Double(pixel.red * 255))
            data[index + 1] = .number(Double(pixel.green * 255))
            data[index + 2] = .number(Double(pixel.blue * 255))
            data[index + 3] = .number(Double(255))
        }
    }
    _ = ctx.putImageData!(imageData, 0, 0)
}

struct ImageView: @unchecked Sendable {
    let width, height: Int
    let buffer: UnsafeMutableBufferPointer<Color>

    subscript(x: Int, y: Int) -> Color {
        get {
            return buffer[y * width + x]
        }
        nonmutating set {
            buffer[y * width + x] = newValue
        }
    }
}

struct Work: Sendable {
    let scene: Scene
    let imageView: ImageView
    let yRange: CountableRange<Int>

    init(scene: Scene, imageView: ImageView, yRange: CountableRange<Int>) {
        self.scene = scene
        self.imageView = imageView
        self.yRange = yRange
    }
    func run() {
        for y in yRange {
            for x in 0..<scene.width {
                let ray = Ray.createPrime(x: x, y: y, scene: scene)
                let color = castRay(scene: scene, ray: ray, depth: 0)
                imageView[x, y] = color
            }
        }
    }
}

func render(scene: Scene, ctx: JSObject, renderTimeElement: JSObject, concurrency: Int, executor: (some TaskExecutor)?) async {

    let imageBuffer = UnsafeMutableBufferPointer<Color>.allocate(capacity: scene.width * scene.height)
    // Initialize the buffer with black color
    imageBuffer.initialize(repeating: .black)
    let imageView = ImageView(width: scene.width, height: scene.height, buffer: imageBuffer)

    let clock = ContinuousClock()
    let start = clock.now

    func updateRenderTime() {
        let renderSceneDuration = clock.now - start
        renderTimeElement.textContent = .string("Render time: \(renderSceneDuration)")
    }

    var checkTimer: JSValue?
    checkTimer = JSObject.global.setInterval!(JSClosure { _ in
        print("Checking thread work...")
        renderInCanvas(ctx: ctx, image: imageView)
        updateRenderTime()
        return .undefined
    }, 250)

    await withTaskGroup(of: Void.self) { group in
        let yStride = scene.height / concurrency
        for i in 0..<concurrency {
            let yRange = i * yStride..<(i + 1) * yStride
            let work = Work(scene: scene, imageView: imageView, yRange: yRange)
            group.addTask(executorPreference: executor) { work.run() }
        }
        // Remaining rows
        if scene.height % concurrency != 0 {
            let work = Work(scene: scene, imageView: imageView, yRange: (concurrency * yStride)..<scene.height)
            group.addTask(executorPreference: executor) { work.run() }
        }
    }

    _ = JSObject.global.clearInterval!(checkTimer!)
    checkTimer = nil

    renderInCanvas(ctx: ctx, image: imageView)
    updateRenderTime()
    imageBuffer.deallocate()
    print("All work done")
}

func onClick() async throws {
    let document = JSObject.global.document

    let canvasElement = document.getElementById("canvas").object!
    let renderTimeElement = document.getElementById("render-time").object!

    let concurrency = max(Int(document.getElementById("concurrency").object!.value.string!) ?? 1, 1)
    let background = document.getElementById("background").object!.checked.boolean!
    let size = Int(document.getElementById("size").object!.value.string ?? "800")!

    let ctx = canvasElement.getContext!("2d").object!

    let scene = createDemoScene(size: size)
    let executor = background ? try await WebWorkerTaskExecutor(numberOfThreads: concurrency) : nil
    canvasElement.width  = .number(Double(scene.width))
    canvasElement.height = .number(Double(scene.height))

    await render(scene: scene, ctx: ctx, renderTimeElement: renderTimeElement, concurrency: concurrency, executor: executor)
    executor?.terminate()
    print("Render done")
}

func main() async throws {
    let renderButtonElement = JSObject.global.document.getElementById("render-button").object!
    let concurrencyElement = JSObject.global.document.getElementById("concurrency").object!
    concurrencyElement.value = JSObject.global.navigator.hardwareConcurrency

    _ = renderButtonElement.addEventListener!("click", JSClosure { _ in
        Task {
            try await onClick()
        }
        return JSValue.undefined
    })
}

Task {
    try await main()
}


#if canImport(wasi_pthread)
import wasi_pthread
import WASILibc

/// Trick to avoid blocking the main thread. pthread_mutex_lock function is used by
/// the Swift concurrency runtime.
@_cdecl("pthread_mutex_lock")
func pthread_mutex_lock(_ mutex: UnsafeMutablePointer<pthread_mutex_t>) -> Int32 {
    // DO NOT BLOCK MAIN THREAD
    var ret: Int32
    repeat {
        ret = pthread_mutex_trylock(mutex)
    } while ret == EBUSY
    return ret
}
#endif
