// Example

import stb_truetype
import Foundation

class STBTextProvider: UITextProvider {
    // TODO: Add a fallback font.
    let font = try! Data(contentsOf: URL(filePath: "/path/to/Arial.ttf"))
    var pixels = [UInt8](repeating: 0, count: 1024 * 1024)
    var chars = [stbtt_bakedchar](repeating: .init(), count: 256)
    
    public init() {
        let atlas = font.withUnsafeBytes { ptr in
            pixels.withUnsafeMutableBufferPointer { pixels in
                chars.withUnsafeMutableBufferPointer { chars in
                    stbtt_BakeFontBitmap(ptr.baseAddress, 0, 32, pixels.baseAddress, 1024, 1024, Int32("\0".utf8CString[0]), 256, chars.baseAddress)
                }
            }
        }
        pixels[0] = .max
    }
    
    public func getQuadForChar(_ char: Character) -> UIRect {
        var x: Float = 0
        var y: Float = 0
        var quad = stbtt_aligned_quad()
        chars.withUnsafeMutableBufferPointer { chars in
            stbtt_GetBakedQuad(chars.baseAddress, 1024, 1024, Int32(char.asciiValue ?? "?".first!.asciiValue!), &x, &y, &quad, 1)
        }
        return UIRect(origin: UIVec2(Double(x), Double(y)), size: UIVec2(Double(quad.x1 - quad.x0), Double(quad.y1 - quad.y0)))
    }
    
    public func getCharInfo(_ char: Character) -> stbtt_bakedchar {
        return chars[Int(char.asciiValue ?? "?".first!.asciiValue!)]
    }
    
    public func measure(text: String, maxWidth: Double?) -> UIVec2 {
        var out = UIVec2(0.0, 0.0)
        var advance = UIVec2(0.0, 32.0)
        for char in text {
            let info = getCharInfo(char)
            advance.x += Double(info.xadvance)
            if let maxWidth = maxWidth, advance.x >= maxWidth {
                out.x = maxWidth
                advance.x = 0
                advance.y += 32
            }
        }
        out.x = max(out.x, advance.x)
        out.y = advance.y
        return out
    }
    
    public func drawText(into list: inout UIDrawList, rect: UIRect, text: String, color: UIColor, fontName: String?, fontSize: Float, alignment: Alignment) {
        var advance = 0.0
        for char in text {
            let info = getCharInfo(char)
            let cw = Double(info.x1 - info.x0)
            let ch = Double(info.y1 - info.y0)
            
            let uvLeft   = Double(info.x0) / 1024.0
            let uvRight  = Double(info.x1) / 1024.0
            let uvTop    = Double(info.y0) / 1024.0
            let uvBottom = Double(info.y1) / 1024.0
            
            list.addQuad(
                bounds: UIRect(
                    origin: rect.origin + (advance + Double(info.xoff), Double(fontSize + info.yoff)),
                    size: UIVec2(cw, ch)
                ),
                uv: UIRect(
                    origin: UIVec2(uvLeft, uvTop),
                    size: UIVec2(uvRight - uvLeft, uvBottom - uvTop)
                ),
                color: color
            )
            advance += Double(info.xadvance)
        }
    }
}

let atlas = STBTextProvider()

let ctx = UIContext(STBTextProvider()) {
    HStack {
        VStack {
            Text("Test 1")
            Text("Test 3")
            Text("Test 4")
        }
        .width(.grow(1))
        .backgroundColor(.white)
        .withStyle(BorderRadiusStyle.self, 100)
        .withStyle(BorderWidthStyle.self, 6)
        
        VStack {
            ScrollView {
                VStack {
                    for i in 1...50 {
                        Text("Item \(i)")
                    }
                }
            }
            .height(.grow(1))
            Text("Test 2")
        }
        .alignment(.trailing)
        .width(.grow(1))
        .backgroundColor(.init(60, 60, 60))
    }
    .width(.relative(1))
    .height(.relative(1))
}

let window = GLFWWindow(width: 1024, height: 720)
let graphics = window.createGraphics(WebGPUGraphics.self)

let uniforms = graphics.createBuffer(label: "Uniform buffer", usage: [.uniform, .copyDest], size: 2048)
let texture = graphics.createTexture(label: "Texture", usage: [.texture, .copyDest], size: [1024, 1024])
let sampler = graphics.createSampler(label: "Sampler")

var rgba = [UInt8](repeating: 0, count: 1024*1024*4)
for i in 0..<1024*1024 {
    let gray = atlas.pixels[i]
    let idx = i * 4
    rgba[idx+0] = 255      // R
    rgba[idx+1] = 255      // G
    rgba[idx+2] = 255      // B
    rgba[idx+3] = gray     // A
}
texture.writeData(rgba)

var projection = glm_mat4()
glm_ortho(0, Float(window.width), Float(window.height), 0, -1, 1, &projection.0)

let verts = graphics.createBuffer(label: "UI vertex buffer", usage: [.vertex, .copyDest], size: 1024 * 1024)
let inds = graphics.createBuffer(label: "UI index buffer", usage: [.index, .copyDest], size: 1024 * 1024)
var drawable = Drawable(passes: ["ui"], vertices: verts, vertexCount: 0, indices: inds, indexCount: 0, uniforms: mat4AsBytes(&projection))

ctx.viewportSize = UIVec2(1024, 720)

let list = ctx.getDrawList()

print("Vertices: \(list.vertices)")

verts.writeData(list.vertices)
drawable.vertexCount = list.vertexCount
inds.writeData(list.indices.withUnsafeBytes { Array($0) })
drawable.indexCount = list.indices.count

ctx.dump()

window.registerOnResize { x, y in
    ctx.viewportSize = UIVec2(Double(x), Double(y))
    glm_ortho(0, Float(window.width), Float(window.height), 0, -1, 1, &projection.0)
    drawable.uniforms = mat4AsBytes(&projection)
}

graphics.registerPass(
    RenderPass(
        name: "ui",
        label: "UI Pass",
        bindings: [
            .buffer(stage: .vertex, type: .uniform, buffer: uniforms),
            .texture(stage: .fragment, sampleType: .float, texture: texture),
            .sampler(stage: .fragment, type: .filtering, sampler: sampler)
        ],
        vertexLayout: [.init(attributes: [.vec2f, .vec2f, .vec4n])],
        vertexShader: Shader(entryPoint: "vs_main", source: """
            struct Uniforms {
                ortho: mat4x4<f32>,
            }
            @group(0) @binding(0) var<uniform> uniforms: Uniforms;
            struct VertexOutput {
                @builtin(position) pos: vec4<f32>,
                @location(0) uv: vec2<f32>,
                @location(1) color: vec4<f32>,
            }
            @vertex
            fn vs_main(@location(0) pos: vec2<f32>, @location(1) uv: vec2<f32>, @location(2) color: vec4<f32>) -> VertexOutput {
                var out: VertexOutput;
                out.pos = uniforms.ortho * vec4<f32>(pos, 0.0, 1.0);
                out.uv = uv;
                out.color = color;
                return out;
            }
            """),
        colorTargets: [ColorTarget(channels: .all, colorBlendOperation: .add, colorBlendSrcFactor: .srcAlpha, colorBlendDestFactor: .oneMinusSrcAlpha, alphaBlendOperation: .add, alphaBlendSrcFactor: .one, alphaBlendDestFactor: .oneMinusSrcAlpha)],
        fragmentShader: Shader(entryPoint: "fs_main", source: """
            @group(0) @binding(1) var texture: texture_2d<f32>;
            @group(0) @binding(2) var sampler_: sampler;
            @fragment
            fn fs_main(@location(0) uv: vec2<f32>, @location(1) color: vec4<f32>) -> @location(0) vec4<f32> {
                return textureSample(texture, sampler_, uv) * color;
            }
            """),
        clearColor: .init(0, 0, 128, .max)
    )
)

graphics.scene.append(drawable)

window.run()