import Foundation
import UIKit

// MARK: - Mesh Transform Structs (matching CAMutableMeshTransform memory layout)

struct MeshTransformMeshFace {
    var indices: (UInt32, UInt32, UInt32, UInt32)
    var w: (Float, Float, Float, Float)
}

struct MeshTransformPoint3D {
    var x: CGFloat
    var y: CGFloat
    var z: CGFloat
}

struct MeshTransformMeshVertex {
    var from: CGPoint
    var to: MeshTransformPoint3D
}

// MARK: - MeshTransform (wraps private CAMutableMeshTransform)

final class GlassMeshTransform {
    typealias Value = NSObject
    typealias Point3D = MeshTransformPoint3D
    typealias Vertex = MeshTransformMeshVertex
    typealias Face = MeshTransformMeshFace

    private var vertices: ContiguousArray<Vertex> = []
    private var faces: ContiguousArray<Face> = []

    func add(_ vertex: Vertex) {
        vertices.append(vertex)
    }

    func add(_ face: Face) {
        faces.append(face)
    }

    func makeValue() -> Value? {
        let result = vertices.withUnsafeMutableBufferPointer { vertices -> NSObject? in
            return faces.withUnsafeMutableBufferPointer { faces -> NSObject? in
                return Self.invokeCreateCustomMethod(
                    vertexCount: UInt(vertices.count),
                    vertices: vertices.baseAddress!,
                    faceCount: UInt(faces.count),
                    faces: faces.baseAddress!,
                    depthNormalization: "none"
                )
            }
        }
        if let result {
            Self.invokeSetSubdivisionSteps(object: result, value: 0)
        }
        return result
    }

    // MARK: - Private API Access

    private static let transformClass: NSObject? = {
        let name = ("CAMutable" as NSString).appendingFormat("MeshTransform")
        return NSClassFromString(name as String) as AnyObject as? NSObject
    }()

    @inline(__always)
    private static func getMethod<T>(object: NSObject, selector: String) -> T? {
        guard let method = object.method(for: NSSelectorFromString(selector)) else { return nil }
        return unsafeBitCast(method, to: T.self)
    }

    private static var cachedCreateCustomMethod: ((@convention(c) (AnyObject, Selector, UInt, UnsafeMutableRawPointer, UInt, UnsafeMutableRawPointer, NSString) -> NSObject?), Selector)?

    private static func invokeCreateCustomMethod(
        vertexCount: UInt,
        vertices: UnsafeMutablePointer<MeshTransformMeshVertex>,
        faceCount: UInt,
        faces: UnsafeMutablePointer<MeshTransformMeshFace>,
        depthNormalization: String
    ) -> NSObject? {
        guard let transformClass else { return nil }
        let rawVertices = UnsafeMutableRawPointer(vertices)
        let rawFaces = UnsafeMutableRawPointer(faces)
        if let cached = cachedCreateCustomMethod {
            return cached.0(transformClass, cached.1, vertexCount, rawVertices, faceCount, rawFaces, depthNormalization as NSString)
        }
        let selectorName = ("meshTransf" as NSString).appending("ormWithVertexCount:vertices:faceCount:faces:depthNormalization:")
        let method: (@convention(c) (AnyObject, Selector, UInt, UnsafeMutableRawPointer, UInt, UnsafeMutableRawPointer, NSString) -> NSObject?)? = getMethod(object: transformClass, selector: selectorName)
        if let method {
            let selector = NSSelectorFromString(selectorName)
            cachedCreateCustomMethod = (method, selector)
            return method(transformClass, selector, vertexCount, rawVertices, faceCount, rawFaces, depthNormalization as NSString)
        }
        return nil
    }

    private static var cachedSetSubdivisionSteps: ((@convention(c) (AnyObject, Selector, Int) -> Void), Selector)?

    private static func invokeSetSubdivisionSteps(object: NSObject, value: Int) {
        if let cached = cachedSetSubdivisionSteps {
            cached.0(object, cached.1, value)
            return
        }
        let method: (@convention(c) (AnyObject, Selector, Int) -> Void)? = getMethod(object: object, selector: "setSubdivisionSteps:")
        if let method {
            let selector = NSSelectorFromString("setSubdivisionSteps:")
            cachedSetSubdivisionSteps = (method, selector)
            method(object, selector, value)
        }
    }
}

// MARK: - Bezier Helpers

private func a(_ a1: CGFloat, _ a2: CGFloat) -> CGFloat {
    return 1.0 - 3.0 * a2 + 3.0 * a1
}

private func b(_ a1: CGFloat, _ a2: CGFloat) -> CGFloat {
    return 3.0 * a2 - 6.0 * a1
}

private func c(_ a1: CGFloat) -> CGFloat {
    return 3.0 * a1
}

private func calcBezier(_ t: CGFloat, _ a1: CGFloat, _ a2: CGFloat) -> CGFloat {
    return ((a(a1, a2) * t + b(a1, a2)) * t + c(a1)) * t
}

private func calcSlope(_ t: CGFloat, _ a1: CGFloat, _ a2: CGFloat) -> CGFloat {
    return 3.0 * a(a1, a2) * t * t + 2.0 * b(a1, a2) * t + c(a1)
}

private func getTForX(_ x: CGFloat, _ x1: CGFloat, _ x2: CGFloat) -> CGFloat {
    var t = x
    for _ in 0..<4 {
        let currentSlope = calcSlope(t, x1, x2)
        if currentSlope == 0.0 {
            return t
        }
        let currentX = calcBezier(t, x1, x2) - x
        t -= currentX / currentSlope
    }
    return t
}

private func bezierPoint(_ x1: CGFloat, _ y1: CGFloat, _ x2: CGFloat, _ y2: CGFloat, _ x: CGFloat) -> CGFloat {
    var value = calcBezier(getTForX(x, x1, x2), y1, y2)
    if value >= 0.997 {
        value = 1.0
    }
    return value
}

// MARK: - Displacement Bezier

struct DisplacementBezier {
    var x1: CGFloat
    var y1: CGFloat
    var x2: CGFloat
    var y2: CGFloat
}

// MARK: - SDF + Displacement

func roundedRectSDF(x: CGFloat, y: CGFloat, width: CGFloat, height: CGFloat, cornerRadius: CGFloat) -> CGFloat {
    let px = x - width / 2
    let py = y - height / 2
    let bx = width / 2
    let by = height / 2
    let qx = abs(px) - bx + cornerRadius
    let qy = abs(py) - by + cornerRadius
    let outsideDist = hypot(max(qx, 0), max(qy, 0))
    let insideDist = min(max(qx, qy), 0)
    return outsideDist + insideDist - cornerRadius
}

private func roundedRectGradient(x: CGFloat, y: CGFloat, width: CGFloat, height: CGFloat, cornerRadius: CGFloat) -> (nx: CGFloat, ny: CGFloat) {
    let px = x - width / 2
    let py = y - height / 2
    let bx = width / 2
    let by = height / 2
    let qx = abs(px) - bx + cornerRadius
    let qy = abs(py) - by + cornerRadius

    var nx: CGFloat = 0
    var ny: CGFloat = 0

    if qx > 0 && qy > 0 {
        let d = hypot(qx, qy)
        if d > 0 {
            nx = qx / d
            ny = qy / d
        }
    } else if qx > qy {
        nx = 1
        ny = 0
    } else {
        nx = 0
        ny = 1
    }

    if px < 0 { nx = -nx }
    if py < 0 { ny = -ny }

    return (nx, ny)
}

private func computeDisplacement(
    x: CGFloat, y: CGFloat,
    width: CGFloat, height: CGFloat,
    cornerRadius: CGFloat,
    edgeDistance: CGFloat,
    bezier: DisplacementBezier
) -> (dx: CGFloat, dy: CGFloat, sdf: CGFloat) {
    let sdf = roundedRectSDF(x: x, y: y, width: width, height: height, cornerRadius: cornerRadius)
    let (nx, ny) = roundedRectGradient(x: x, y: y, width: width, height: height, cornerRadius: cornerRadius)

    let inwardX = -nx
    let inwardY = -ny
    let distFromEdge = -sdf
    let weight = max(0, min(1, 1.0 - distFromEdge / edgeDistance))

    var dx = inwardX * weight
    var dy = inwardY * weight

    let mag = hypot(dx, dy)
    if mag > 0 {
        let newMag = bezierPoint(bezier.x1, bezier.y1, bezier.x2, bezier.y2, mag)
        let scale = newMag / mag
        dx *= scale
        dy *= scale
    }

    return (dx, dy, sdf)
}

// MARK: - Mesh Template Cache

private struct GlassMeshCacheKey: Hashable {
    var cornerRadius: CGFloat
    var edgeDistance: CGFloat
    var cornerResolution: Int
    var outerEdgeDistance: CGFloat
    var bezierX1: CGFloat
    var bezierY1: CGFloat
    var bezierX2: CGFloat
    var bezierY2: CGFloat

    init(cornerRadius: CGFloat, edgeDistance: CGFloat, cornerResolution: Int, outerEdgeDistance: CGFloat, bezier: DisplacementBezier) {
        self.cornerRadius = cornerRadius
        self.edgeDistance = edgeDistance
        self.cornerResolution = cornerResolution
        self.outerEdgeDistance = outerEdgeDistance
        self.bezierX1 = bezier.x1
        self.bezierY1 = bezier.y1
        self.bezierX2 = bezier.x2
        self.bezierY2 = bezier.y2
    }
}

private struct GlassMeshTemplate {
    struct VertexTemplate {
        var baseX: CGFloat
        var sizeScaleX: CGFloat
        var baseY: CGFloat
        var sizeScaleY: CGFloat
        var dispX: CGFloat
        var dispY: CGFloat
        var depth: CGFloat
    }

    var vertices: ContiguousArray<VertexTemplate>
    var faces: ContiguousArray<GlassMeshTransform.Face>
}

private var glassMeshTemplateCache: [GlassMeshCacheKey: GlassMeshTemplate] = [:]

private func instantiateGlassMesh(
    from template: GlassMeshTemplate,
    size: CGSize,
    displacementMagnitudeU: CGFloat,
    displacementMagnitudeV: CGFloat
) -> GlassMeshTransform {
    let W = size.width
    let H = size.height
    let insetPoints: CGFloat = -1.0
    let insetUOffset = insetPoints / W
    let insetVOffset = insetPoints / H
    let usableUNorm = (W - insetPoints * 2) / W
    let usableVNorm = (H - insetPoints * 2) / H

    let transform = GlassMeshTransform()
    for v in template.vertices {
        let worldX = v.baseX + v.sizeScaleX * W
        let worldY = v.baseY + v.sizeScaleY * H
        let u = worldX / W
        let vCoord = worldY / H
        let mappedU = insetUOffset + u * usableUNorm
        let mappedV = insetVOffset + vCoord * usableVNorm
        let fromX = max(0.0, min(1.0, mappedU + v.dispX * displacementMagnitudeU))
        let fromY = max(0.0, min(1.0, mappedV + v.dispY * displacementMagnitudeV))
        transform.add(GlassMeshTransform.Vertex(
            from: CGPoint(x: fromX, y: fromY),
            to: GlassMeshTransform.Point3D(x: mappedU, y: mappedV, z: v.depth)
        ))
    }
    for face in template.faces {
        transform.add(face)
    }
    return transform
}

private func generateGlassMeshTemplate(
    cornerRadius: CGFloat,
    edgeDistance: CGFloat,
    cornerResolution: Int,
    outerEdgeDistance: CGFloat,
    bezier: DisplacementBezier
) -> GlassMeshTemplate {
    let clampedRadius = cornerRadius
    let refW = max(4 * clampedRadius, 100)
    let refH = max(4 * clampedRadius, 100)

    var vertices = ContiguousArray<GlassMeshTemplate.VertexTemplate>()
    var faces = ContiguousArray<GlassMeshTransform.Face>()
    var vertexIndex: Int = 0

    func templateDisplacement(worldX: CGFloat, worldY: CGFloat) -> (CGFloat, CGFloat) {
        let (rawDispX, rawDispY, sdf) = computeDisplacement(
            x: worldX, y: worldY,
            width: refW, height: refH,
            cornerRadius: clampedRadius,
            edgeDistance: edgeDistance,
            bezier: bezier
        )
        let distToEdge = max(0.0, -sdf)
        let edgeBand = max(0.0, outerEdgeDistance)
        let edgeBoost: CGFloat
        if edgeBand > 0 {
            let t = max(0.0, min(1.0, (edgeBand - distToEdge) / edgeBand))
            edgeBoost = 1.0 + t * t * (3 - 2 * t) * 0.5
        } else {
            edgeBoost = 1.0
        }
        return (rawDispX * edgeBoost, rawDispY * edgeBoost)
    }

    func addVertex(baseX: CGFloat, scaleX: CGFloat, baseY: CGFloat, scaleY: CGFloat, depth: CGFloat = 0) -> Int {
        let worldX = baseX + scaleX * refW
        let worldY = baseY + scaleY * refH
        let (dispX, dispY) = templateDisplacement(worldX: worldX, worldY: worldY)
        vertices.append(GlassMeshTemplate.VertexTemplate(
            baseX: baseX, sizeScaleX: scaleX,
            baseY: baseY, sizeScaleY: scaleY,
            dispX: dispX, dispY: dispY, depth: depth
        ))
        let idx = vertexIndex
        vertexIndex += 1
        return idx
    }

    func addQuadFace(_ i0: Int, _ i1: Int, _ i2: Int, _ i3: Int) {
        faces.append(GlassMeshTransform.Face(
            indices: (UInt32(i0), UInt32(i1), UInt32(i2), UInt32(i3)),
            w: (0.0, 0.0, 0.0, 0.0)
        ))
    }

    let angularStepsBase = max(3, cornerResolution)
    let angularSteps = angularStepsBase % 2 == 0 ? angularStepsBase : angularStepsBase + 1
    let radialSteps = max(2, cornerResolution)
    let horizontalSegments = max(2, cornerResolution / 2 + 1)
    let verticalSegments = max(2, cornerResolution / 2 + 1)
    let R = clampedRadius

    func depthFactorsWithOuterBand(count: Int, band: CGFloat, maxRadius: CGFloat) -> [CGFloat] {
        guard count > 0, maxRadius > 0 else { return [0, 1] }
        let bandNorm = max(0, min(1, band / maxRadius))
        let innerSegments = max(1, count - 1)
        let innerMax = max(0, 1 - bandNorm)
        var factors: [CGFloat] = (0...innerSegments).map { i in
            innerMax * CGFloat(i) / CGFloat(innerSegments)
        }
        func appendUnique(_ value: CGFloat) {
            if let last = factors.last, abs(last - value) < 1e-4 { return }
            factors.append(value)
        }
        appendUnique(innerMax)
        appendUnique(1.0)
        return factors
    }

    let depthFactors = depthFactorsWithOuterBand(count: radialSteps, band: outerEdgeDistance, maxRadius: R)
    let angularFactors = (0...angularSteps).map { CGFloat($0) / CGFloat(angularSteps) }
    let outerToInner = depthFactors.reversed()

    let topXCoeffs: [(base: CGFloat, scale: CGFloat)] = (0...horizontalSegments).map { i in
        let t = CGFloat(i) / CGFloat(horizontalSegments)
        return (base: R * (1 - 2 * t), scale: t)
    }
    let sideYCoeffs: [(base: CGFloat, scale: CGFloat)] = (0...verticalSegments).map { j in
        let t = CGFloat(j) / CGFloat(verticalSegments)
        return (base: R * (1 - 2 * t), scale: t)
    }
    let topYCoeffs: [(base: CGFloat, scale: CGFloat)] = outerToInner.map { factor in
        (base: R * (1 - factor), scale: 0)
    }
    let bottomYCoeffs: [(base: CGFloat, scale: CGFloat)] = depthFactors.map { factor in
        (base: -R * (1 - factor), scale: 1)
    }
    let leftXCoeffs: [(base: CGFloat, scale: CGFloat)] = outerToInner.map { factor in
        (base: R * (1 - factor), scale: 0)
    }
    let rightXCoeffs: [(base: CGFloat, scale: CGFloat)] = depthFactors.map { factor in
        (base: -R * (1 - factor), scale: 1)
    }

    func buildGridTemplate(
        xCoeffs: [(base: CGFloat, scale: CGFloat)],
        yCoeffs: [(base: CGFloat, scale: CGFloat)]
    ) {
        var indexGrid: [[Int]] = []
        for yc in yCoeffs {
            var row: [Int] = []
            for xc in xCoeffs {
                row.append(addVertex(baseX: xc.base, scaleX: xc.scale, baseY: yc.base, scaleY: yc.scale))
            }
            indexGrid.append(row)
        }
        let numRows = indexGrid.count - 1
        let numCols = indexGrid.first!.count - 1
        for row in 0..<numRows {
            for col in 0..<numCols {
                addQuadFace(
                    indexGrid[row][col],
                    indexGrid[row][col + 1],
                    indexGrid[row + 1][col + 1],
                    indexGrid[row + 1][col]
                )
            }
        }
    }

    func buildCornerTemplate(
        centerBaseX: CGFloat, centerScaleX: CGFloat,
        centerBaseY: CGFloat, centerScaleY: CGFloat,
        startAngle: CGFloat, endAngle: CGFloat
    ) {
        let ringRadials = outerToInner.filter { $0 > 0 }
        guard !ringRadials.isEmpty else { return }

        var ringIndices: [[Int]] = []
        for radial in ringRadials {
            let r = R * radial
            var row: [Int] = []
            for t in angularFactors {
                let angle = startAngle + (endAngle - startAngle) * t
                let offsetX = r * cos(angle)
                let offsetY = r * sin(angle)
                row.append(addVertex(
                    baseX: centerBaseX + offsetX, scaleX: centerScaleX,
                    baseY: centerBaseY + offsetY, scaleY: centerScaleY
                ))
            }
            ringIndices.append(row)
        }

        for r in 0..<(ringIndices.count - 1) {
            let outerRing = ringIndices[r]
            let innerRing = ringIndices[r + 1]
            for i in 0..<(outerRing.count - 1) {
                addQuadFace(outerRing[i], outerRing[i + 1], innerRing[i + 1], innerRing[i])
            }
        }

        if let innermostRing = ringIndices.last {
            let ringSegments = innermostRing.count - 1
            guard ringSegments >= 2 else { return }

            let centerAnchor = addVertex(
                baseX: centerBaseX, scaleX: centerScaleX,
                baseY: centerBaseY, scaleY: centerScaleY,
                depth: -0.02
            )
            let stride = 2
            var i = 0
            while i + 2 <= ringSegments {
                addQuadFace(centerAnchor, innermostRing[i], innermostRing[i + 1], innermostRing[i + 2])
                i += stride
            }
            if i < ringSegments {
                addQuadFace(centerAnchor, innermostRing[ringSegments - 1], innermostRing[ringSegments], innermostRing[ringSegments])
            }
        }
    }

    // Edge strips
    buildGridTemplate(xCoeffs: topXCoeffs, yCoeffs: topYCoeffs)
    buildGridTemplate(xCoeffs: topXCoeffs, yCoeffs: bottomYCoeffs)
    buildGridTemplate(xCoeffs: leftXCoeffs, yCoeffs: sideYCoeffs)
    buildGridTemplate(xCoeffs: rightXCoeffs, yCoeffs: sideYCoeffs)

    // Center patch
    buildGridTemplate(xCoeffs: topXCoeffs, yCoeffs: sideYCoeffs)

    // Corners
    buildCornerTemplate(centerBaseX: R, centerScaleX: 0, centerBaseY: R, centerScaleY: 0, startAngle: .pi, endAngle: 1.5 * .pi)
    buildCornerTemplate(centerBaseX: -R, centerScaleX: 1, centerBaseY: R, centerScaleY: 0, startAngle: 1.5 * .pi, endAngle: 2 * .pi)
    buildCornerTemplate(centerBaseX: -R, centerScaleX: 1, centerBaseY: -R, centerScaleY: 1, startAngle: .pi / 2, endAngle: 0)
    buildCornerTemplate(centerBaseX: R, centerScaleX: 0, centerBaseY: -R, centerScaleY: 1, startAngle: .pi, endAngle: .pi / 2)

    return GlassMeshTemplate(vertices: vertices, faces: faces)
}

// MARK: - Public Entry Point

func generateGlassMesh(
    size: CGSize,
    cornerRadius: CGFloat,
    edgeDistance: CGFloat,
    displacementMagnitudeU: CGFloat,
    displacementMagnitudeV: CGFloat,
    cornerResolution: Int,
    outerEdgeDistance: CGFloat,
    bezier: DisplacementBezier
) -> GlassMeshTransform {
    let clampedRadius = min(cornerRadius, min(size.width, size.height) / 2)

    let key = GlassMeshCacheKey(
        cornerRadius: clampedRadius,
        edgeDistance: edgeDistance,
        cornerResolution: cornerResolution,
        outerEdgeDistance: outerEdgeDistance,
        bezier: bezier
    )
    let template: GlassMeshTemplate
    if let cached = glassMeshTemplateCache[key] {
        template = cached
    } else {
        template = generateGlassMeshTemplate(
            cornerRadius: clampedRadius,
            edgeDistance: edgeDistance,
            cornerResolution: cornerResolution,
            outerEdgeDistance: outerEdgeDistance,
            bezier: bezier
        )
        glassMeshTemplateCache[key] = template
    }
    return instantiateGlassMesh(
        from: template,
        size: size,
        displacementMagnitudeU: displacementMagnitudeU,
        displacementMagnitudeV: displacementMagnitudeV
    )
}
