import Foundation
import simd
import Synchronization
import SwiftUI

struct ReferencePointStorage {
    var position: SIMD2<Float>
    var orbit: [SIMD2<Float>]
    var period: Int32
    var maxIter: Int32
}

final class ReferencePoint: Sendable {
    let storage: Mutex<ReferencePointStorage>

    init(position: SIMD2<Float>, scale: Float) {
        storage = .init(.init(position: position, orbit: Array(repeating: SIMD2<Float>(0, 0), count: 1024), period: 0, maxIter: 100))
        Task {
            await calculateOrbit(scale: scale)
        }
    }

    func updatePosition(position: SIMD2<Float>, scale: Float) {
        storage.withLock {
            $0.position = position
            Task {
                await calculateOrbit(scale: scale)
            }
        }
    }

    func calculateOrbit(scale: Float) async {
        var z = SIMD2<Float>(0, 0)
        storage.withLock {
            $0.maxIter = Int32(min(100 + log2(Float(scale)) * 25, 1000))

            for i in 0..<1024 {
                $0.orbit[i] = z

                let real = z.x * z.x - z.y * z.y + $0.position.x
                let imag = 2 * z.x * z.y + $0.position.y
                z = SIMD2<Float>(real, imag)

                if (z.x * z.x + z.y * z.y) > 4 {
                    $0.maxIter = Int32(i)
                    break
                }

                if i > 20 {
                    for j in 1...20 {
                        if abs(z.x - $0.orbit[i-j].x) < 1e-6 && abs(z.y - $0.orbit[i-j].y) < 1e-6 {
                            $0.period = Int32(j)
                            $0.maxIter = Int32(i)
                            return
                        }
                    }
                }
            }
        }
    }

    func getData() -> sending Data {
        var data = Data(capacity: MemoryLayout<SIMD2<Float>>.size * 1025 + MemoryLayout<Int32>.size * 2)

        storage.withLock {
            var positionCopy = $0.position
            var orbitCopy = $0.orbit
            var periodCopy = $0.period
            var maxIterCopy = $0.maxIter
            data.append(Data(bytes: &positionCopy, count: MemoryLayout<SIMD2<Float>>.size))
            data.append(Data(bytes: &orbitCopy, count: MemoryLayout<SIMD2<Float>>.size * 1024))
            data.append(Data(bytes: &periodCopy, count: MemoryLayout<Int32>.size))
            data.append(Data(bytes: &maxIterCopy, count: MemoryLayout<Int32>.size))
        }
        return data
    }

    func getPosition() -> SIMD2<Float> {
        storage.withLock {
            return $0.position
        }
    }
}

extension View {
    func mandelbrotShader(offset: CGSize, scale: CGFloat, color: Color) -> some View {
        modifier(MandelbrotShader(offset: offset, scale: scale, color: color))
    }
}

struct MandelbrotShader: ViewModifier {
    let offset: CGSize
    let scale: CGFloat
    let color: Color

    @State private var referencePoint: ReferencePoint

    init(offset: CGSize, scale: CGFloat, color: Color) {
        self.offset = offset
        self.scale = scale
        self.color = color
        self.referencePoint = ReferencePoint(position: SIMD2<Float>(-0.5, 0), scale: Float(scale))
    }

    func body(content: Content) -> some View {
        content
            .visualEffect { [referencePoint] content, proxy in
                let components = color.resolve(in: EnvironmentValues())

                let existingPos = referencePoint.getPosition()
                let currentPos = SIMD2<Float>(
                    Float(-0.5 + offset.width),
                    Float(offset.height)
                )
                if simd_distance(currentPos, existingPos) > 0.1 / Float(scale) {
                    referencePoint.updatePosition(position: currentPos, scale: Float(scale)) // UNSURE IF THIS WILL MUTATE THE @STATE VERSION

                    // AVOID THE NEED FOR MUTATION HERE BY PASSING getData() TO .data() BELOW
                    //                        self.referenceData = referencePoint?.getData() ?? Data()
                    //                        print(self.referenceData)
                }

                return content
                    .colorEffect(ShaderLibrary.mandelbrot(
                        .float2(proxy.size),
                        .float2(Float(offset.width), Float(offset.height)),
                        .float(Float(scale)),
                        .float3(Float(components.red), Float(components.green), Float(components.blue)),
                        .data(referencePoint.getData())
                    ))
            }
    }
}