import Foundation
import FoundationModels
import NaturalLanguage
import SwiftRs

struct SendableCallback: @unchecked Sendable {
    let callback: @convention(c) (UnsafePointer<UInt8>?, Int, UnsafeMutableRawPointer) -> Void
    let sender: UnsafeMutableRawPointer
}

struct SendableStreamCallback: @unchecked Sendable {
    let callback: @convention(c) (UnsafePointer<UInt8>?, Int, Bool, UnsafeMutableRawPointer) -> Void
    let sender: UnsafeMutableRawPointer
}

struct SendableIntCallback: @unchecked Sendable {
    let callback: @convention(c) (Int32, UnsafeMutableRawPointer) -> Void
    let sender: UnsafeMutableRawPointer
}

struct SendableFloatArrayCallback: @unchecked Sendable {
    let callback: @convention(c) (UnsafePointer<Float>?, Int, UnsafeMutableRawPointer) -> Void
    let sender: UnsafeMutableRawPointer
}

@available(macOS 26.0, *)
internal func createOptions(temperature: Float, maxTokens: Int32) -> GenerationOptions {
    var options = GenerationOptions()
    if temperature >= 0 {
        options.temperature = Double(temperature)
    }
    if maxTokens > 0 {
        options.maximumResponseTokens = Int(maxTokens)
    }
    return options
}

@available(macOS 26.0, *)
actor SessionManager {
    static let shared = SessionManager()
    private var sessions: [String: LanguageModelSession] = [:]
    private var sessionInstructions: [String: String] = [:]
    
    func getSession(id: String, instructions: String?) -> LanguageModelSession {
        let newInstructions = instructions ?? ""
        
        // If session exists and instructions match, return it
        if let existing = sessions[id], sessionInstructions[id] == newInstructions {
            return existing
        }
        
        // Otherwise, create/recreate the session
        let newSession: LanguageModelSession
        if !newInstructions.isEmpty {
            newSession = LanguageModelSession(instructions: newInstructions)
        } else {
            newSession = LanguageModelSession()
        }
        
        sessions[id] = newSession
        sessionInstructions[id] = newInstructions
        return newSession
    }
    
    func removeSession(id: String) {
        sessions.removeValue(forKey: id)
        sessionInstructions.removeValue(forKey: id)
    }
}

@available(macOS 26.0, *)
@_cdecl("call_apple_llm")
public func callAppleLLM(
    prompt: UnsafePointer<SRString>, 
    instructions: UnsafePointer<SRString>,
    sessionId: UnsafePointer<SRString>,
    temperature: Float, 
    maxTokens: Int32, 
    sender: UnsafeMutableRawPointer, 
    callback: @escaping @convention(c) (UnsafePointer<UInt8>?, Int, UnsafeMutableRawPointer) -> Void
) {
    let promptStr = prompt.pointee.toString()
    let instructionsStr = instructions.pointee.toString()
    let sessionIdStr = sessionId.pointee.toString()
    let container = SendableCallback(callback: callback, sender: sender)
    let nativeOptions = createOptions(temperature: temperature, maxTokens: maxTokens)
    
    Task.detached {
        do {
            let model = SystemLanguageModel.default
            
            guard model.availability == .available else {
                let errorMsg = "Error: Apple Intelligence is not available (\(model.availability))"
                let data = errorMsg.data(using: .utf8)!
                data.withUnsafeBytes { b in container.callback(b.baseAddress?.assumingMemoryBound(to: UInt8.self), data.count, container.sender) }
                return
            }

            let session = await SessionManager.shared.getSession(id: sessionIdStr, instructions: instructionsStr)
            
            let response = try await session.respond(to: promptStr, options: nativeOptions)
            
            let data = response.content.data(using: .utf8)!
            data.withUnsafeBytes { buffer in
                if let baseAddress = buffer.baseAddress {
                    let ptr = baseAddress.assumingMemoryBound(to: UInt8.self)
                    container.callback(ptr, data.count, container.sender)
                }
            }
        } catch {
            let errorMsg = "Error: \(error.localizedDescription)"
            let data = errorMsg.data(using: .utf8)!
            data.withUnsafeBytes { buffer in
                if let baseAddress = buffer.baseAddress {
                    let ptr = baseAddress.assumingMemoryBound(to: UInt8.self)
                    container.callback(ptr, data.count, container.sender)
                }
            }
        }
    }
}

@available(macOS 26.0, *)
@_cdecl("call_apple_llm_stream")
public func callAppleLLMStream(
    prompt: UnsafePointer<SRString>, 
    instructions: UnsafePointer<SRString>,
    sessionId: UnsafePointer<SRString>,
    temperature: Float, 
    maxTokens: Int32, 
    sender: UnsafeMutableRawPointer, 
    callback: @escaping @convention(c) (UnsafePointer<UInt8>?, Int, Bool, UnsafeMutableRawPointer) -> Void
) {
    let promptStr = prompt.pointee.toString()
    let instructionsStr = instructions.pointee.toString()
    let sessionIdStr = sessionId.pointee.toString()
    let container = SendableStreamCallback(callback: callback, sender: sender)
    let nativeOptions = createOptions(temperature: temperature, maxTokens: maxTokens)
    
    Task.detached {
        do {
            let model = SystemLanguageModel.default
            guard model.availability == .available else {
                let errorMsg = "Error: AI not available (\(model.availability))"
                let data = errorMsg.data(using: .utf8)!
                data.withUnsafeBytes { b in container.callback(b.baseAddress?.assumingMemoryBound(to: UInt8.self), data.count, true, container.sender) }
                return
            }

            let session = await SessionManager.shared.getSession(id: sessionIdStr, instructions: instructionsStr)
            let stream = session.streamResponse(to: promptStr, options: nativeOptions)
            
            for try await partial in stream {
                let data = partial.content.data(using: .utf8)!
                data.withUnsafeBytes { buffer in
                    if let baseAddress = buffer.baseAddress {
                        let ptr = baseAddress.assumingMemoryBound(to: UInt8.self)
                        container.callback(ptr, data.count, false, container.sender)
                    }
                }
            }
            container.callback(nil, 0, true, container.sender)
        } catch {
            let errorMsg = "Error: \(error.localizedDescription)"
            let data = errorMsg.data(using: .utf8)!
            data.withUnsafeBytes { buffer in
                if let baseAddress = buffer.baseAddress {
                    let ptr = baseAddress.assumingMemoryBound(to: UInt8.self)
                    container.callback(ptr, data.count, true, container.sender)
                }
            }
        }
    }
}

@available(macOS 26.0, *)
@_cdecl("apple_llm_prewarm")
public func appleLLMPrewarm() {
    Task.detached {
        let session = LanguageModelSession()
        session.prewarm()
    }
}

@available(macOS 26.0, *)
@_cdecl("apple_llm_clear_session")
public func appleLLMClearSession(sessionId: UnsafePointer<SRString>) {
    let id = sessionId.pointee.toString()
    Task.detached {
        await SessionManager.shared.removeSession(id: id)
    }
}

@available(macOS 26.0, *)
@_cdecl("apple_llm_count_tokens")
public func appleLLMCountTokens(text: UnsafePointer<SRString>, sender: UnsafeMutableRawPointer, callback: @escaping @convention(c) (Int32, UnsafeMutableRawPointer) -> Void) {
    let textStr = text.pointee.toString()
    let container = SendableIntCallback(callback: callback, sender: sender)
    Task.detached {
        let count = Int32(textStr.count / 4)
        container.callback(count, container.sender)
    }
}

@available(macOS 26.0, *)
@_cdecl("apple_llm_get_context_size")
public func appleLLMGetContextSize(sender: UnsafeMutableRawPointer, callback: @escaping @convention(c) (Int32, UnsafeMutableRawPointer) -> Void) {
    let container = SendableIntCallback(callback: callback, sender: sender)
    Task.detached {
        container.callback(4096, container.sender)
    }
}

@_cdecl("apple_nl_get_embedding")
public func appleNLGetEmbedding(text: UnsafePointer<SRString>, sender: UnsafeMutableRawPointer, callback: @escaping @convention(c) (UnsafePointer<Float>?, Int, UnsafeMutableRawPointer) -> Void) {
    let textStr = text.pointee.toString()
    let container = SendableFloatArrayCallback(callback: callback, sender: sender)
    Task.detached {
        if let embedding = NLEmbedding.sentenceEmbedding(for: .english), let vector = embedding.vector(for: textStr) {
            let floats = vector.map { Float($0) }
            floats.withUnsafeBufferPointer { buffer in
                container.callback(buffer.baseAddress, floats.count, container.sender)
            }
        } else {
            container.callback(nil, 0, container.sender)
        }
    }
}
