import _CJavaScriptKit
#if hasFeature(Embedded) && os(WASI)
import _Concurrency
#endif

/// `JSClosureProtocol` wraps Swift closure objects for use in JavaScript. Conforming types
/// are responsible for managing the lifetime of the closure they wrap, but can delegate that
/// task to the user by requiring an explicit `release()` call.
public protocol JSClosureProtocol: JSValueCompatible {

    /// Release this function resource.
    /// After calling `release`, calling this function from JavaScript will fail.
    func release()
}

/// `JSOneshotClosure` is a JavaScript function that can be called only once. This class can be used
/// for optimized memory management when compared to the common `JSClosure`.
public class JSOneshotClosure: JSObject, JSClosureProtocol {
    private var hostFuncRef: JavaScriptHostFuncRef = 0

    public init(
        file: String = #fileID,
        line: UInt32 = #line,
        _ body: @escaping (sending [JSValue]) -> JSValue
    ) {
        // 1. Fill `id` as zero at first to access `self` to get `ObjectIdentifier`.
        super.init(id: 0)

        // 2. Create a new JavaScript function which calls the given Swift function.
        hostFuncRef = JavaScriptHostFuncRef(bitPattern: ObjectIdentifier(self))
        _id = withExtendedLifetime(JSString(file)) { file in
            swjs_create_oneshot_function(hostFuncRef, line, file.asInternalJSRef())
        }

        // 3. Retain the given body in static storage by `funcRef`.
        JSClosure.sharedClosures.wrappedValue[hostFuncRef] = .init(
            object: self,
            body: {
                defer { self.release() }
                return body($0)
            },
            fileID: file,
            line: line
        )
    }

    @available(*, unavailable, message: "JSOneshotClosure does not support dictionary literal initialization")
    public required init(dictionaryLiteral elements: (String, JSValue)...) {
        fatalError("JSOneshotClosure does not support dictionary literal initialization")
    }

    #if compiler(>=5.5) && (!hasFeature(Embedded) || os(WASI))
    /// Creates a new `JSOneshotClosure` that calls the given Swift function asynchronously.
    ///
    /// - Parameters:
    ///   - priority: The priority of the new unstructured Task created under the hood.
    ///   - body: The Swift function to call asynchronously.
    ///   - file: The file identifier where the function is defined.
    ///   - line: The line number where the function is defined.
    @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
    public static func async(
        priority: TaskPriority? = nil,
        file: String = #fileID,
        line: UInt32 = #line,
        _ body: @escaping (sending [JSValue]) async throws(JSException) -> JSValue
    ) -> JSOneshotClosure {
        JSOneshotClosure(file: file, line: line, makeAsyncClosure(priority: priority, body))
    }

    /// Creates a new `JSOneshotClosure` that calls the given Swift function asynchronously.
    ///
    /// - Parameters:
    ///   - taskExecutor: The executor preference of the new unstructured Task created under the hood.
    ///   - priority: The priority of the new unstructured Task created under the hood.
    ///   - body: The Swift function to call asynchronously.
    ///   - file: The file identifier where the function is defined.
    ///   - line: The line number where the function is defined.
    @available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
    public static func async(
        executorPreference taskExecutor: (any TaskExecutor)? = nil,
        priority: TaskPriority? = nil,
        file: String = #fileID,
        line: UInt32 = #line,
        _ body: @escaping (sending [JSValue]) async throws(JSException) -> JSValue
    ) -> JSOneshotClosure {
        JSOneshotClosure(
            file: file,
            line: line,
            makeAsyncClosure(executorPreference: taskExecutor, priority: priority, body)
        )
    }
    #endif

    /// Release this function resource.
    /// After calling `release`, calling this function from JavaScript will fail.
    public func release() {
        JSClosure.sharedClosures.wrappedValue[hostFuncRef] = nil
    }
}

/// `JSClosure` represents a JavaScript function the body of which is written in Swift.
/// This type can be passed as a callback handler to JavaScript functions.
///
/// e.g.
/// ```swift
/// let eventListener = JSClosure { _ in
///     ...
///     return JSValue.undefined
/// }
///
/// button.addEventListener!("click", JSValue.function(eventListener))
/// ...
/// button.removeEventListener!("click", JSValue.function(eventListener))
/// ```
///
public class JSClosure: JSObject, JSClosureProtocol {

    class SharedJSClosure {
        // Note: 6.0 compilers built with assertions enabled crash when calling
        // `removeValue(forKey:)` on a dictionary with value type containing
        // `sending`. Wrap the value type with a struct to avoid the crash.
        struct Entry {
            let object: JSObject
            let body: (sending [JSValue]) -> JSValue
            #if Tracing
            let fileID: String
            let line: UInt32
            #endif

            init(object: JSObject, body: @escaping (sending [JSValue]) -> JSValue, fileID: String, line: UInt32) {
                self.object = object
                self.body = body
                #if Tracing
                self.fileID = fileID
                self.line = line
                #endif
            }
        }
        private var storage: [JavaScriptHostFuncRef: Entry] = [:]
        init() {}

        subscript(_ key: JavaScriptHostFuncRef) -> Entry? {
            get { storage[key] }
            set { storage[key] = newValue }
        }
    }

    // Note: Retain the closure object itself also to avoid funcRef conflicts
    fileprivate static let sharedClosures = LazyThreadLocal {
        SharedJSClosure()
    }

    private var hostFuncRef: JavaScriptHostFuncRef = 0

    #if JAVASCRIPTKIT_WITHOUT_WEAKREFS
    private var isReleased: Bool = false
    #endif

    @available(
        *,
        deprecated,
        message:
            "This initializer will be removed in the next minor version update. Please use `init(_ body: @escaping ([JSValue]) -> JSValue)` and add `return .undefined` to the end of your closure"
    )
    @_disfavoredOverload
    public convenience init(_ body: @escaping ([JSValue]) -> Void) {
        self.init({
            body($0)
            return .undefined
        })
    }

    public init(
        file: String = #fileID,
        line: UInt32 = #line,
        _ body: @escaping (sending [JSValue]) -> JSValue
    ) {
        // 1. Fill `id` as zero at first to access `self` to get `ObjectIdentifier`.
        super.init(id: 0)

        // 2. Create a new JavaScript function which calls the given Swift function.
        hostFuncRef = JavaScriptHostFuncRef(bitPattern: ObjectIdentifier(self))
        _id = withExtendedLifetime(JSString(file)) { file in
            swjs_create_function(hostFuncRef, line, file.asInternalJSRef())
        }

        // 3. Retain the given body in static storage by `funcRef`.
        Self.sharedClosures.wrappedValue[hostFuncRef] = .init(
            object: self,
            body: body,
            fileID: file,
            line: line
        )
    }

    @available(*, unavailable, message: "JSClosure does not support dictionary literal initialization")
    public required init(dictionaryLiteral elements: (String, JSValue)...) {
        fatalError("JSClosure does not support dictionary literal initialization")
    }

    #if compiler(>=5.5) && (!hasFeature(Embedded) || os(WASI))
    /// Creates a new `JSClosure` that calls the given Swift function asynchronously.
    ///
    /// - Parameters:
    ///   - priority: The priority of the new unstructured Task created under the hood.
    ///   - body: The Swift function to call asynchronously.
    ///   - file: The file identifier where the function is defined.
    ///   - line: The line number where the function is defined.
    @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
    public static func async(
        priority: TaskPriority? = nil,
        file: String = #fileID,
        line: UInt32 = #line,
        _ body: @escaping @isolated(any) (sending [JSValue]) async throws(JSException) -> JSValue
    ) -> JSClosure {
        JSClosure(file: file, line: line, makeAsyncClosure(priority: priority, body))
    }

    /// Creates a new `JSClosure` that calls the given Swift function asynchronously.
    ///
    /// - Parameters:
    ///   - taskExecutor: The executor preference of the new unstructured Task created under the hood.
    ///   - priority: The priority of the new unstructured Task created under the hood.
    ///   - body: The Swift function to call asynchronously.
    ///   - file: The file identifier where the function is defined.
    ///   - line: The line number where the function is defined.
    @available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
    public static func async(
        executorPreference taskExecutor: (any TaskExecutor)? = nil,
        priority: TaskPriority? = nil,
        file: String = #fileID,
        line: UInt32 = #line,
        _ body: @escaping (sending [JSValue]) async throws(JSException) -> JSValue
    ) -> JSClosure {
        JSClosure(file: file, line: line, makeAsyncClosure(executorPreference: taskExecutor, priority: priority, body))
    }
    #endif

    #if JAVASCRIPTKIT_WITHOUT_WEAKREFS
    deinit {
        guard isReleased else {
            fatalError("release() must be called on JSClosure objects manually before they are deallocated")
        }
    }
    #endif
}

/// Heap storage for typed closures generated by BridgeJS.
///
/// The storage is owned by the JavaScript closure and keeps the call-site metadata
/// so that calls after `release()` can report a helpful error.
@_spi(BridgeJS) public final class _BridgeJSTypedClosureBox<Signature>: _BridgedSwiftClosureBox {
    @_spi(BridgeJS) public var closure: Signature
    var isManuallyReleased: Bool = false

    init(_ closure: Signature) {
        self.closure = closure
    }
}

/// A typed handle to a JavaScript closure generated by BridgeJS.
///
/// Use `JSTypedClosure` when you need to pass a Swift closure to JavaScript and then
/// manage its lifetime explicitly from Swift. You must call `release()` when the closure
/// is no longer needed. After releasing, any attempt to invoke the closure from JavaScript
/// throws a `JSException` (no crash), including a file/line message pointing to where the
/// typed closure was created.
///
/// The underlying JavaScript function is also registered with a `FinalizationRegistry` as a
/// safety net to release the Swift closure storage if `release()` is forgotten, but this
/// is non-deterministic and should not be relied on for timely cleanup.
///
/// ```swift
/// @JSFunction func someJSFunction(_ c: JSTypedClosure<(Int) -> Int>) throws(JSException)
/// let closure = JSTypedClosure<(Int) -> Int> { $0 + 1 }
/// defer { closure.release() }
/// try someJSFunction(closure)
/// ```
public struct JSTypedClosure<Signature> {
    private let storage: _BridgeJSTypedClosureBox<Signature>
    private let _jsObject: JSObject

    @_spi(BridgeJS)
    public init(
        makeClosure: (_ pointer: UnsafeMutableRawPointer, _ fileID: UnsafePointer<UInt8>, _ line: UInt32) -> Int32,
        body: Signature,
        fileID: StaticString,
        line: UInt32
    ) {
        let storage = _BridgeJSTypedClosureBox<Signature>(body)
        // NOTE: Perform an unbalanced retain to give the closure box storage ownership
        // to the JavaScript side. It will be balanced by two code paths:
        //
        // 1. When the JS closure made by `makeClosure` is GC'ed, the FinalizationRegistry
        //    will call `_bjs_release_swift_closure` to release the closure box storage.
        // 2. When Swift side manually calls `release()` on the closure, it will release
        //    the closure box storage.
        let pointer = Unmanaged.passRetained(storage).toOpaque()
        let id = makeClosure(pointer, fileID.utf8Start, line)
        self._jsObject = JSObject(id: UInt32(bitPattern: id))
        self.storage = storage
    }

    @_spi(BridgeJS)
    public func bridgeJSLowerParameter() -> Int32 {
        return _jsObject.bridgeJSLowerParameter()
    }

    @_spi(BridgeJS)
    public func bridgeJSLowerReturn() -> Int32 {
        return _jsObject.bridgeJSLowerReturn()
    }

    /// The underlying JavaScript function handle.
    ///
    /// Use this when you need to store or compare the JS function identity on the JavaScript side.
    public var jsObject: JSObject { _jsObject }

    /// Releases the Swift closure storage owned by this typed closure.
    ///
    /// Call this exactly once when the closure is no longer needed by JavaScript.
    /// After release, any attempt to invoke the closure from JavaScript throws a `JSException`
    /// describing a call on a released `JSTypedClosure`.
    /// Calling `release()` multiple times has no effect.
    public func release() {
        guard !storage.isManuallyReleased else {
            // Already manually released
            return
        }
        storage.isManuallyReleased = true
        // Unregister the closure from the FinalizationRegistry because we here decrement
        // the refcount of the storage, so the finalizer won't be responsible for releasing it.
        _swift_js_closure_unregister(Int32(bitPattern: _jsObject.id))
        // Balance the retain given to the JavaScript side
        Unmanaged.passUnretained(storage).release()
    }
}

#if compiler(>=5.5) && (!hasFeature(Embedded) || os(WASI))
@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
private func makeAsyncClosure(
    priority: TaskPriority?,
    _ body: @escaping @isolated(any) (sending [JSValue]) async throws(JSException) -> JSValue
) -> ((sending [JSValue]) -> JSValue) {
    { arguments in
        JSPromise { resolver in
            // NOTE: The context is fully transferred to the unstructured task
            // isolation but the compiler can't prove it yet, so we need to
            // use `@unchecked Sendable` to make it compile with the Swift 6 mode.
            struct Context: @unchecked Sendable {
                let resolver: (JSPromise.Result) -> Void
                let arguments: [JSValue]
                let body: (sending [JSValue]) async throws(JSException) -> JSValue
            }
            let context = Context(resolver: resolver, arguments: arguments, body: body)
            Task(priority: priority) {
                do throws(JSException) {
                    let result = try await context.body(context.arguments)
                    context.resolver(.success(result))
                } catch {
                    context.resolver(.failure(error.thrownValue))
                }
            }
        }.jsValue()
    }
}

@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
private func makeAsyncClosure(
    executorPreference taskExecutor: (any TaskExecutor)?,
    priority: TaskPriority?,
    _ body: @escaping (sending [JSValue]) async throws(JSException) -> JSValue
) -> ((sending [JSValue]) -> JSValue) {
    { arguments in
        JSPromise { resolver in
            // NOTE: The context is fully transferred to the unstructured task
            // isolation but the compiler can't prove it yet, so we need to
            // use `@unchecked Sendable` to make it compile with the Swift 6 mode.
            struct Context: @unchecked Sendable {
                let resolver: (JSPromise.Result) -> Void
                let arguments: [JSValue]
                let body: (sending [JSValue]) async throws(JSException) -> JSValue
            }
            let context = Context(resolver: resolver, arguments: arguments, body: body)
            Task(executorPreference: taskExecutor, priority: priority) {
                do throws(JSException) {
                    let result = try await context.body(context.arguments)
                    context.resolver(.success(result))
                } catch {
                    context.resolver(.failure(error.thrownValue))
                }
            }
        }.jsValue()
    }
}
#endif

// MARK: - `JSClosure` mechanism note
//
// 1. Create a thunk in the JavaScript world, which has a reference
//    to a Swift closure.
// ┌─────────────────────┬──────────────────────────┐
// │     Swift side      │   JavaScript side        │
// │                     │                          │
// │                     │                          │
// │                     │   ┌──────[Thunk]───────┐ │
// │          ┌ ─ ─ ─ ─ ─│─ ─│─ ─ ─ ─ ─ ┐         │ │
// │          ↓          │   │          │         │ │
// │  [Swift Closure]    │   │  Host Function ID  │ │
// │                     │   │                    │ │
// │                     │   └────────────────────┘ │
// └─────────────────────┴──────────────────────────┘
//
// 2. When the thunk function is invoked, it calls the Swift closure via
//    `_call_host_function` and receives the result through a callback.
// ┌─────────────────────┬──────────────────────────┐
// │     Swift side      │   JavaScript side        │
// │                     │                          │
// │                     │                          │
// │                   Apply ┌──────[Thunk]───────┐ │
// │          ┌ ─ ─ ─ ─ ─│─ ─│─ ─ ─ ─ ─ ┐         │ │
// │          ↓          │   │          │         │ │
// │  [Swift Closure]    │   │  Host Function ID  │ │
// │          │          │   │                    │ │
// │          │          │   └────────────────────┘ │
// │          │          │                    ↑     │
// │          │        Apply                  │     │
// │          └─[Result]─┼───>[Callback func]─┘     │
// │                     │                          │
// └─────────────────────┴──────────────────────────┘

/// Returns true if the host function has been already released, otherwise false.
@_cdecl("_call_host_function_impl")
func _call_host_function_impl(
    _ hostFuncRef: JavaScriptHostFuncRef,
    _ argv: UnsafePointer<RawJSValue>,
    _ argc: Int32,
    _ callbackFuncRef: JavaScriptObjectRef
) -> Bool {
    guard let entry = JSClosure.sharedClosures.wrappedValue[hostFuncRef] else {
        return true
    }
    #if Tracing
    let traceEnd = JSTracingHooks.beginJSClosureCall(
        JSTracing.JSClosureCallInfo(fileID: entry.fileID, line: UInt(entry.line))
    )
    #endif
    var arguments: [JSValue] = []
    for i in 0..<Int(argc) {
        arguments.append(argv[i].jsValue)
    }
    #if Tracing
    defer { traceEnd?() }
    #endif
    let result = entry.body(arguments)
    let callbackFuncRef = JSObject(id: callbackFuncRef)
    _ = callbackFuncRef(result)
    return false
}

/// [WeakRefs](https://github.com/tc39/proposal-weakrefs) are already Stage 4,
/// but was added recently enough that older browser versions don’t support it.
/// Build with `-Xswiftc -DJAVASCRIPTKIT_WITHOUT_WEAKREFS` to disable the relevant behavior.
#if JAVASCRIPTKIT_WITHOUT_WEAKREFS

// MARK: - Legacy Closure Types

extension JSClosure {
    public func release() {
        isReleased = true
        Self.sharedClosures.wrappedValue[hostFuncRef] = nil
    }
}

@_cdecl("_free_host_function_impl")
func _free_host_function_impl(_ hostFuncRef: JavaScriptHostFuncRef) {}

#else

extension JSClosure {

    @available(*, deprecated, message: "JSClosure.release() is no longer necessary")
    public func release() {}

}

@_cdecl("_free_host_function_impl")
func _free_host_function_impl(_ hostFuncRef: JavaScriptHostFuncRef) {
    JSClosure.sharedClosures.wrappedValue[hostFuncRef] = nil
}
#endif

#if compiler(>=6.0) && hasFeature(Embedded)
// cdecls currently don't work in embedded, and expose for wasm only works >=6.0
@_expose(wasm, "swjs_call_host_function")
public func _swjs_call_host_function(
    _ hostFuncRef: JavaScriptHostFuncRef,
    _ argv: UnsafePointer<RawJSValue>,
    _ argc: Int32,
    _ callbackFuncRef: JavaScriptObjectRef
) -> Bool {

    _call_host_function_impl(hostFuncRef, argv, argc, callbackFuncRef)
}

@_expose(wasm, "swjs_free_host_function")
public func _swjs_free_host_function(_ hostFuncRef: JavaScriptHostFuncRef) {
    _free_host_function_impl(hostFuncRef)
}
#endif
