#pragma once

#include "core/slang-blob.h"
#include "core/slang-dictionary.h"
#include "core/slang-list.h"
#include "core/slang-memory-arena.h"
#include "replay-shared.h"
#include "replay-stream.h"

#include <cstdint>
#include <cstring>
#include <limits>
#include <mutex>
#include <slang-com-helper.h>
#include <slang-com-ptr.h>
#include <slang.h>
#include <type_traits>

namespace SlangRecord
{

using Slang::Dictionary;
using Slang::List;
using Slang::MemoryArena;
using Slang::String;

// Fwd declare from proxy-base.h (non-template, work with canonical ISlangUnknown*)
SLANG_API ISlangUnknown* wrapObject(ISlangUnknown* obj);
SLANG_API ISlangUnknown* unwrapObject(ISlangUnknown* proxy);

/// Type-safe wrapObject: canonicalizes the input via queryInterface, wraps it,
/// then QIs the result back to T. Safe for diamond-inheritance types.
template<typename T>
inline T* wrapObject(T* obj)
{
    ISlangUnknown* wrapped = wrapObject(toSlangUnknown(obj));
    return wrapped ? toSlangInterface<T>(wrapped) : nullptr;
}

/// Type-safe unwrapObject: canonicalizes the input via queryInterface, unwraps it,
/// then QIs the result back to T. Safe for diamond-inheritance types.
template<typename T>
inline T* unwrapObject(T* obj)
{
    ISlangUnknown* impl = unwrapObject(toSlangUnknown(obj));
    return impl ? toSlangInterface<T>(impl) : nullptr;
}

/// Handle constants for interface tracking.
/// Handles 0-255 are reserved for special meanings.
constexpr uint64_t kNullHandle = 0; ///< Null pointer
constexpr uint64_t kCustomFileSystemHandle =
    2; ///< User-provided custom file system (that has not yet been registered with a handle)
constexpr uint64_t kDefaultFileSystemHandle =
    3; ///< Default file system (when user doesn't provide one)
constexpr uint64_t kFirstValidHandle = 0x100; ///< First handle for tracked objects

/// Replay file sanity limits. They bound attacker-controlled allocation sizes
/// before deserializers allocate replay-owned buffers.
constexpr uint32_t kMaxReplayStringLength = 64 * 1024 * 1024;
constexpr uint64_t kMaxReplayArrayCount = 16 * 1024 * 1024;
constexpr size_t kMaxReplayArrayAllocationSize = 64 * 1024 * 1024;
constexpr size_t kMaxReplayTotalAllocationSize = 256 * 1024 * 1024;
constexpr size_t kMaxReplayAllocationToStreamSizeRatio = 16;

/// Fixed-size index entry for the call index stream.
/// Each entry stores only the byte offset of a call in stream.bin.
/// The signature and this-handle are read from the main stream on demand.
#pragma pack(push, 1)
struct CallIndexEntry
{
    uint64_t streamPosition; ///< Byte offset in stream.bin where call begins

    /// Total size of this struct (must be fixed for file I/O)
    static constexpr size_t kSize = sizeof(uint64_t);
};
#pragma pack(pop)

static_assert(
    sizeof(CallIndexEntry) == CallIndexEntry::kSize,
    "CallIndexEntry must have expected fixed size");

// Replay exceptions can be thrown across the slang shared-library boundary
// and caught by type in the test executable when slang-test runs in-process
// (server-count=1). They therefore carry SLANG_EXCEPTION_TYPE_VISIBLE
// (core/slang-exception.h), which exports the RTTI on Itanium-ABI
// (non-Windows) targets so the typed catch matches across dynamic libraries;
// see that header for the full explanation.

/// Exception thrown when trying to record an untracked interface.
class SLANG_EXCEPTION_TYPE_VISIBLE UntrackedInterfaceException : public Slang::Exception
{
public:
    UntrackedInterfaceException(ISlangUnknown* obj)
        : Slang::Exception("Attempted to record untracked interface"), m_object(obj)
    {
    }
    ISlangUnknown* getObject() const { return m_object; }

private:
    ISlangUnknown* m_object;
};

/// Exception thrown when a handle is not found during playback.
class SLANG_EXCEPTION_TYPE_VISIBLE HandleNotFoundException : public Slang::Exception
{
public:
    HandleNotFoundException(uint64_t handle)
        : Slang::Exception(String("Handle not found: ") + String(handle)), m_handle(handle)
    {
    }
    uint64_t getHandle() const { return m_handle; }

private:
    uint64_t m_handle;
};

/// Exception thrown when a TypeReflection cannot be resolved during playback.
class SLANG_EXCEPTION_TYPE_VISIBLE UnresolvedTypeException : public Slang::Exception
{
public:
    UnresolvedTypeException(slang::TypeReflection* type)
        : Slang::Exception(String("Unresolved type: ") + String(type ? type->getName() : "<null>"))
        , m_type(type)
    {
    }
    slang::TypeReflection* getType() const { return m_type; }

private:
    slang::TypeReflection* m_type;
};

/// Operating mode for the replay system.
enum class Mode : uint8_t
{
    Idle,     ///< No data captured, operations are no-ops
    Record,   ///< Writing data to a stream
    Sync,     ///< Writing data and comparing to reference stream for determinism verification
    Playback, ///< Reading data from a stream
};

/// Flags indicating the role of a value being recorded.
/// Used to determine replay verification behavior.
enum class RecordFlag : uint8_t
{
    None = 0,               ///< No special handling
    Input = 1 << 0,         ///< Function input argument (verify on replay)
    Output = 1 << 1,        ///< Function output argument (capture on replay)
    InOut = Input | Output, ///< Input/output argument
    ReturnValue = 1 << 2,   ///< Function return value (capture on replay)
    ThisPtr = 1 << 3,       ///< 'this' pointer for method calls
};

inline RecordFlag operator|(RecordFlag a, RecordFlag b)
{
    return static_cast<RecordFlag>(static_cast<uint8_t>(a) | static_cast<uint8_t>(b));
}

inline RecordFlag operator&(RecordFlag a, RecordFlag b)
{
    return static_cast<RecordFlag>(static_cast<uint8_t>(a) & static_cast<uint8_t>(b));
}

inline bool hasFlag(RecordFlag flags, RecordFlag flag)
{
    return (static_cast<uint8_t>(flags) & static_cast<uint8_t>(flag)) != 0;
}

/// Type identifiers for serialized values.
enum class TypeId : uint8_t
{
    Int8 = 0x01,
    Int16 = 0x02,
    Int32 = 0x03,
    Int64 = 0x04,
    UInt8 = 0x05,
    UInt16 = 0x06,
    UInt32 = 0x07,
    UInt64 = 0x08,
    Float32 = 0x09,
    Float64 = 0x0A,
    Bool = 0x0B,
    String = 0x10,
    Blob = 0x11,
    Array = 0x12,
    ObjectHandle = 0x13,
    Null = 0x14,
    TypeReflectionRef = 0x15, ///< TypeReflection reference (module handle + type name)
    Error = 0xEE,             ///< Error marker - indicates an exception occurred
};

SLANG_API const char* getTypeIdName(TypeId id);

/// Exception thrown when type mismatch occurs during deserialization.
class SLANG_EXCEPTION_TYPE_VISIBLE TypeMismatchException : public Slang::Exception
{
public:
    SLANG_API TypeMismatchException(TypeId expected, TypeId actual);
    TypeId getExpected() const { return m_expected; }
    TypeId getActual() const { return m_actual; }

private:
    TypeId m_expected, m_actual;
};

/// Exception thrown when data mismatch occurs during sync mode verification.
class SLANG_EXCEPTION_TYPE_VISIBLE DataMismatchException : public Slang::Exception
{
public:
    SLANG_API DataMismatchException(size_t offset, size_t size);
    size_t getOffset() const { return m_offset; }
    size_t getSize() const { return m_size; }

private:
    size_t m_offset, m_size;
};

/// Require the current playback stream position to contain a byte range.
inline void requireReplayStreamBytes(const ReplayStream& stream, size_t offset, size_t size)
{
    const size_t streamPosition = stream.getPosition();
    const size_t streamSize = stream.getSize();
    if (streamPosition > streamSize || size > streamSize - streamPosition)
        throw DataMismatchException(offset, size);
}

/// Validate a replayed array count before allocation and element deserialization.
inline void validateReplayArrayCount(
    const ReplayStream& stream,
    uint64_t arrayCount,
    uint64_t maxCountForCountType,
    size_t elementSize,
    size_t minStreamBytesPerElement,
    size_t countOffset)
{
    if (arrayCount > maxCountForCountType || arrayCount > kMaxReplayArrayCount)
        throw DataMismatchException(countOffset, sizeof(arrayCount));

    if (arrayCount > uint64_t((std::numeric_limits<size_t>::max)()))
        throw DataMismatchException(countOffset, sizeof(arrayCount));

    const size_t sizeCount = static_cast<size_t>(arrayCount);
    if (elementSize == 0 || minStreamBytesPerElement == 0 ||
        sizeCount > kMaxReplayArrayAllocationSize / elementSize)
        throw DataMismatchException(countOffset, sizeof(arrayCount));

    const size_t streamPosition = stream.getPosition();
    const size_t streamSize = stream.getSize();
    if (streamPosition > streamSize)
        throw DataMismatchException(streamPosition, sizeof(arrayCount));

    const size_t bytesAvailable = streamSize - streamPosition;
    if (arrayCount > uint64_t(bytesAvailable / minStreamBytesPerElement))
        throw DataMismatchException(streamPosition, sizeof(arrayCount));
}

template<typename T>
inline constexpr size_t getReplayArrayMinStreamBytesPerElement()
{
    using ElementType = typename std::remove_cv<T>::type;
    return (std::is_arithmetic<ElementType>::value || std::is_enum<ElementType>::value) ? sizeof(T)
                                                                                        : 1;
}

/// Unified serializer for binary I/O during record/replay.
/// Provides a uniform API for both reading and writing serialized data.
/// Owns its own ReplayStream and MemoryArena internally.
///
/// The context operates in one of four modes:
/// - Idle: No operations performed (default when env var not set)
/// - Record: Writing data to a stream
/// - Sync: Writing data and comparing to reference stream for determinism
/// - Playback: Reading data from a stream
class ReplayContext
{
public:
    /// Get the global singleton instance (for recording).
    /// Thread-safe. The singleton starts in Idle or Record mode based on env var.
    SLANG_API static ReplayContext& get();

    /// Return the singleton if it has already been constructed, or nullptr.
    /// Thread-safe - uses a mutex to synchronize with get().
    /// Use this in cleanup paths (e.g. slang_shutdown) to avoid constructing
    /// the singleton just to tear it down.
    SLANG_API static ReplayContext* tryGet();

    /// Destroy the singleton instance if it exists, freeing all resources.
    /// Thread-safe - uses a mutex to synchronize with get() and tryGet().
    /// Safe to call when no singleton exists (no-op). After this call,
    /// get() will lazily re-create a fresh instance if called again.
    SLANG_API static void destroySingleton();

    /// Create an idle context.
    /// Will switch to Record mode if SLANG_RECORD_LAYER=1 is set.
    SLANG_API ReplayContext();

    /// Create a playback context from existing data.
    SLANG_API ReplayContext(const void* data, size_t size);

    /// Create a sync context that records while verifying against reference data.
    SLANG_API ReplayContext(const void* referenceData, size_t referenceSize, bool syncMode);

    /// Destructor - must be in DLL to properly free Dictionary memory.
    SLANG_API ~ReplayContext();

    /// Get the current operating mode.
    SLANG_API Mode getMode() const { return m_mode; }

    /// Check if the context is active (not Idle).
    SLANG_API bool isActive() const { return m_mode != Mode::Idle; }

    /// Set the operating mode.
    /// When entering Record mode, sets up mirror file for crash-safe capture.
    /// When leaving Record mode, closes the mirror file.
    SLANG_API void setMode(Mode mode);

    /// Convenience methods for common mode checks.
    SLANG_API bool isIdle() const { return m_mode == Mode::Idle; }
    SLANG_API bool isRecording() const { return m_mode == Mode::Record; }
    SLANG_API bool isSyncing() const { return m_mode == Mode::Sync; }
    SLANG_API bool isPlayback() const { return m_mode == Mode::Playback; }
    SLANG_API bool isReading() const { return m_mode == Mode::Playback; }
    SLANG_API bool isWriting() const { return m_mode == Mode::Record || m_mode == Mode::Sync; }

    /// Enable recording (sets mode to Record if currently Idle).
    SLANG_API void enable();

    /// Disable recording (sets mode to Idle).
    SLANG_API void disable();

    /// Reset the context to initial state (clears streams and arena, mode becomes Idle).
    SLANG_API void reset();

    /// Switch to play mode, clearing all local state and resetting stream to 0, but
    /// keeping the stream data present so we can replay what's just happened
    SLANG_API void switchToPlayback();

    /// Switch to sync mode, copying recorded data to reference stream and resetting
    /// for verification pass
    SLANG_API void switchToSync();

    // =========================================================================
    // TTY Logging (Live Trace)
    // =========================================================================

    /// Enable or disable live TTY logging of recorded calls.
    /// When enabled, prints call signatures to stderr as they are recorded.
    /// Can also be enabled via SLANG_RECORD_LOG=1 environment variable.
    SLANG_API void setTtyLogging(bool enable);

    /// Check if TTY logging is enabled.
    SLANG_API bool isTtyLogging() const { return m_ttyLogging; }

    /// Record an error marker in the stream.
    /// Call this when an exception occurs to mark the error location in the stream.
    SLANG_API void recordError(const char* message);

    // =========================================================================
    // Replay Directory Management
    // =========================================================================

    /// Set the base directory for replays (default: ".slang-replays").
    /// Must be called before enabling recording.
    /// Note: If SLANG_RECORD_PATH environment variable is set, it overrides
    /// the base directory and uses the exact path specified.
    SLANG_API void setReplayDirectory(const char* path);

    /// Get the current replay base directory.
    SLANG_API const char* getReplayDirectory() const;

    /// Get the path to the current recording session folder.
    /// Returns nullptr if not currently recording.
    SLANG_API const char* getCurrentReplayPath() const;

    /// Load a replay from a folder path (reads stream.bin).
    /// Switches to Playback mode on success.
    SLANG_API SlangResult loadReplay(const char* folderPath);

    /// Load the most recent replay from the replay directory.
    /// Switches to Playback mode on success.
    SLANG_API SlangResult loadLatestReplay();

    /// Find the most recent replay folder in the given directory.
    /// Returns empty string if no replays found.
    SLANG_API static String findLatestReplayFolder(const char* baseDir);

    // =========================================================================
    // Recording
    // =========================================================================

    SLANG_API ReplayStream& getStream() { return m_stream; }
    SLANG_API const ReplayStream& getStream() const { return m_stream; }
    SLANG_API MemoryArena& getArena() { return m_arena; }
    SLANG_API size_t testsOnlyGetReplayArenaAllocationSize() const
    {
        return m_replayArenaAllocationSize;
    }
    SLANG_API void testsOnlyRequireReplayArenaAllocation(size_t offset, size_t size)
    {
        requireReplayArenaAllocation(offset, size);
    }

    /// Allocates replay-owned arena storage from the module that owns the replay context.
    SLANG_API void* allocateReplayArena(size_t sizeInBytes, size_t alignment);

    /// Lock the context for thread-safe access.
    /// Returns an RAII lock guard.
    SLANG_API std::unique_lock<std::recursive_mutex> lock()
    {
        return std::unique_lock<std::recursive_mutex>(m_mutex);
    }

    // Basic types
    SLANG_API void record(RecordFlag flags, int8_t& value);
    SLANG_API void record(RecordFlag flags, int16_t& value);
    SLANG_API void record(RecordFlag flags, int32_t& value);
    SLANG_API void record(RecordFlag flags, int64_t& value);
    SLANG_API void record(RecordFlag flags, uint8_t& value);
    SLANG_API void record(RecordFlag flags, uint16_t& value);
    SLANG_API void record(RecordFlag flags, uint32_t& value);
    SLANG_API void record(RecordFlag flags, uint64_t& value);
    SLANG_API void record(RecordFlag flags, float& value);
    SLANG_API void record(RecordFlag flags, double& value);
    SLANG_API void record(RecordFlag flags, bool& value);
    SLANG_API void record(RecordFlag flags, const char*& str);

    // Blob by hash - hashes content, stores to disk, records hash in stream.
    // During playback, loads content from disk by hash.
    // This is the primary mechanism for recording ISlangBlob values.
    SLANG_API void recordBlobByHash(RecordFlag flags, ISlangBlob*& blob);

    // Arrays with count - calls record() on each element
    template<typename T, typename CountT>
    void recordArray(RecordFlag flags, T*& arr, CountT& count);
    template<typename T, typename CountT>
    void recordArray(RecordFlag flags, const T*& arr, CountT& count);

    /// Parse a function signature to extract "ClassName::methodName" format.
    /// Works with both MSVC __FUNCSIG__ and GCC/Clang __PRETTY_FUNCTION__.
    /// Returns the normalized signature, or the original if parsing fails.
    SLANG_API static const char* parseSignature(
        const char* signature,
        char* buffer,
        size_t bufferSize);

    /// Begin recording a method call.
    /// Records the function signature and 'this' pointer as a tracked handle.
    /// Also writes an index entry to the index stream for quick navigation.
    template<typename T>
    inline void beginCall(const char* signature, T* thisPtr)
    {
        ensureInitialized();
        if (!isActive())
            return;
        // Parse and record the normalized signature
        char normalizedSig[256];
        const char* parsed = parseSignature(signature, normalizedSig, sizeof(normalizedSig));

        // Log to TTY if enabled
        if (m_ttyLogging)
            logCall(parsed, thisPtr);

        // Write index entry before recording to main stream (for correct position)
        // Get the canonical ISlangUnknown* identity for this proxy
        ISlangUnknown* thisUnknown = toSlangUnknown(thisPtr);

        if (isWriting())
        {
            writeIndexEntry();
        }

        // Note: the parsed signature is fixed/known, so can be treated as a verifiable output.
        record(RecordFlag::Output, parsed);
        recordInterfaceImpl<ISlangUnknown>(RecordFlag::Input, thisUnknown);
    }

    /// Begin recording a static/free function call.
    /// Records only the function signature.
    /// Also writes an index entry to the index stream for quick navigation.
    void beginStaticCall(const char* signature)
    {
        ensureInitialized();
        if (!isActive())
            return;
        char normalizedSig[256];
        const char* parsed = parseSignature(signature, normalizedSig, sizeof(normalizedSig));

        // Log to TTY if enabled
        if (m_ttyLogging)
            logCall(parsed, nullptr);

        // Write index entry before recording to main stream (for correct position)
        if (isWriting())
            writeIndexEntry();

        record(RecordFlag::Input, parsed);
        uint64_t nh = kNullHandle;
        recordHandle(RecordFlag::Input, nh);
    }

    /// Insert a user-defined marker into the replay stream.
    /// During recording, writes a labeled marker entry.
    /// During playback, reads and logs the marker (no-op otherwise).
    void marker(const char* label)
    {
        ensureInitialized();
        if (!isActive())
            return;
        const char* sig = "__marker__";
        record(RecordFlag::Input, sig);
        uint64_t nh = kNullHandle;
        recordHandle(RecordFlag::Input, nh);
        record(RecordFlag::Input, label);
    }

    // Enum types - record as int32_t
    template<typename EnumT>
    void recordEnum(RecordFlag flags, EnumT& value);

    // Slang enum types
    SLANG_API void record(RecordFlag flags, SlangSeverity& value);
    SLANG_API void record(RecordFlag flags, SlangParameterCategory& value);
    SLANG_API void record(RecordFlag flags, SlangBindableResourceType& value);
    SLANG_API void record(RecordFlag flags, SlangCompileTarget& value);
    SLANG_API void record(RecordFlag flags, SlangContainerFormat& value);
    SLANG_API void record(RecordFlag flags, SlangPassThrough& value);
    SLANG_API void record(RecordFlag flags, SlangArchiveType& value);
    SLANG_API void record(RecordFlag flags, SlangFloatingPointMode& value);
    SLANG_API void record(RecordFlag flags, SlangFpDenormalMode& value);
    SLANG_API void record(RecordFlag flags, SlangLineDirectiveMode& value);
    SLANG_API void record(RecordFlag flags, SlangSourceLanguage& value);
    SLANG_API void record(RecordFlag flags, SlangProfileID& value);
    SLANG_API void record(RecordFlag flags, SlangCapabilityID& value);
    SLANG_API void record(RecordFlag flags, SlangMatrixLayoutMode& value);
    SLANG_API void record(RecordFlag flags, SlangStage& value);
    SLANG_API void record(RecordFlag flags, SlangDebugInfoLevel& value);
    SLANG_API void record(RecordFlag flags, SlangDebugInfoFormat& value);
    SLANG_API void record(RecordFlag flags, SlangOptimizationLevel& value);
    SLANG_API void record(RecordFlag flags, SlangPathType& value);
    SLANG_API void record(RecordFlag flags, OSPathKind& value);
    SLANG_API void record(RecordFlag flags, SlangEmitSpirvMethod& value);
    SLANG_API void record(RecordFlag flags, slang::LayoutRules& value);
    SLANG_API void record(RecordFlag flags, slang::CompilerOptionName& value);
    SLANG_API void record(RecordFlag flags, slang::CompilerOptionValueKind& value);
    SLANG_API void record(RecordFlag flags, slang::ContainerType& value);
    SLANG_API void record(RecordFlag flags, slang::SpecializationArg::Kind& value);
    SLANG_API void record(RecordFlag flags, SlangLanguageVersion& value);
    SLANG_API void record(RecordFlag flags, slang::BuiltinModuleName& value);

    // POD and complex structs
    SLANG_API void record(RecordFlag flags, SlangUUID& value);
    SLANG_API void record(RecordFlag flags, slang::CompilerOptionValue& value);
    SLANG_API void record(RecordFlag flags, slang::CompilerOptionEntry& value);
    SLANG_API void record(RecordFlag flags, slang::PreprocessorMacroDesc& value);
    SLANG_API void record(RecordFlag flags, slang::TargetDesc& value);
    SLANG_API void record(RecordFlag flags, slang::SessionDesc& value);
    SLANG_API void record(RecordFlag flags, slang::SpecializationArg& value);
    SLANG_API void record(RecordFlag flags, SlangGlobalSessionDesc& value);

    // Blob interface - serialized by content hash, not tracked as COM interface
    SLANG_API void record(RecordFlag flags, ISlangBlob*& obj);

    // COM interface pointers - handle tracking is done internally
    SLANG_API void record(RecordFlag flags, ISlangFileSystem*& obj);
    SLANG_API void record(RecordFlag flags, ISlangFileSystemExt*& obj);
    SLANG_API void record(RecordFlag flags, ISlangMutableFileSystem*& obj);
    SLANG_API void record(RecordFlag flags, ISlangSharedLibrary*& obj);
    SLANG_API void record(RecordFlag flags, slang::IGlobalSession*& obj);
    SLANG_API void record(RecordFlag flags, slang::ISession*& obj);
    SLANG_API void record(RecordFlag flags, slang::IModule*& obj);
    SLANG_API void record(RecordFlag flags, slang::IComponentType*& obj);
    SLANG_API void record(RecordFlag flags, slang::IEntryPoint*& obj);
    SLANG_API void record(RecordFlag flags, slang::ITypeConformance*& obj);
    SLANG_API void record(RecordFlag flags, slang::ICompileRequest*& obj);

    // TypeReflection - stored as module handle + type name
    // Unlike COM interfaces, TypeReflection cannot be wrapped, so we identify
    // them by their owning module and full type name.
    SLANG_API void record(RecordFlag flags, slang::TypeReflection*& type);

    // Object handles (COM interface pointers mapped to IDs)
    // Public for testing purposes
    SLANG_API void recordHandle(RecordFlag flags, uint64_t& handleId);

    // ==========================================================================
    // Proxy <-> Implementation Mapping
    // ==========================================================================

    /// Register a proxy-implementation pair.
    /// Call this when wrapping an implementation with a proxy.
    template<typename ProxyT, typename ImplT>
    inline uint64_t registerProxy(ProxyT* proxy, ImplT* implementation)
    {
        ISlangUnknown* proxyUnknown = toSlangUnknown(proxy);
        ISlangUnknown* implUnknown = toSlangUnknown(implementation);
        return registerProxyImpl(proxyUnknown, implUnknown);
    }

    /// Unregister a proxy when it is destroyed.
    /// Call this from the proxy destructor to clean up the mappings.
    /// Accepts a pre-computed ISlangUnknown* identity (required because the
    /// destructor can't safely call queryInterface — ref count is already 0).
    inline void unregisterProxy(ISlangUnknown* proxyIdentity)
    {
        unregisterProxyImpl(proxyIdentity);
    }

    /// Registers `obj` in the handle registry and returns its handle, for tests.
    ///
    /// A production proxy is registered by wrapObject() as part of being
    /// created; a test that needs a proxy the orphan bookkeeping will treat as
    /// live has to register it explicitly, because
    /// releaseOrphanedPlaybackProxies() skips anything absent from the registry.
    template<typename ProxyT>
    inline uint64_t testsOnlyRegisterProxy(ProxyT* obj)
    {
        ISlangUnknown* objUnknown = toSlangUnknown(obj);
        return testOnlyRegisterProxyImpl(objUnknown);
    }

    /// Notes an orphaned playback reference for `proxy`, for tests.
    ///
    /// The production caller is the replay dispatcher, which a unit test cannot
    /// drive, so without these wrappers the bookkeeping is only observable as a
    /// sanitizer leak or double-free on a path no test reaches.
    inline void testsOnlyNoteOrphanedProxy(ISlangUnknown* proxy)
    {
        notePlaybackOrphanedProxy(proxy);
    }

    /// Returns how many orphaned playback references are noted for `proxy`, for
    /// tests. Zero when the proxy is not tracked at all.
    inline uint32_t testsOnlyGetOrphanedRefCount(ISlangUnknown* proxy) const
    {
        return testOnlyGetOrphanedRefCountImpl(proxy);
    }

    /// Runs the playback-teardown sweep that releases every noted orphaned
    /// reference, for tests. Production runs it from reset(), switchTo*(), and
    /// destroySingleton().
    inline void testsOnlyReleaseOrphanedProxies() { releaseOrphanedPlaybackProxies(); }

    /// Get or create a proxy for an implementation.
    /// If a proxy already exists, returns it. Otherwise returns nullptr.
    template<typename ImplT>
    inline ISlangUnknown* getProxy(ImplT* implementation)
    {
        if (implementation == nullptr)
            return nullptr;
        ISlangUnknown* implUnknown = toSlangUnknown(implementation);
        return getProxyImpl(implUnknown);
    }

    /// Get the implementation
    template<typename ProxyT>
    inline ISlangUnknown* getImplementation(ProxyT* proxy)
    {
        if (proxy == nullptr)
            return nullptr;
        ISlangUnknown* proxyUnknown = toSlangUnknown(proxy);
        return getImplementationImpl(proxyUnknown);
    }

    /// Get the next handle value that will be assigned.
    /// Useful for testing to know how many objects have been registered.
    SLANG_API uint64_t getNextHandle() const { return m_nextHandle; }

    /// Check if an object is registered.
    template<typename ProxyT>
    inline bool isInterfaceRegistered(ProxyT* obj) const
    {
        ISlangUnknown* objUnknown = toSlangUnknown(obj);
        return isInterfaceRegisteredImpl(objUnknown);
    }

    /// Get handle for an object (throws if not registered).
    template<typename ProxyT>
    inline uint64_t getProxyHandle(ProxyT* obj) const
    {
        ISlangUnknown* objUnknown = toSlangUnknown(obj);
        return getProxyHandleImpl(objUnknown);
    }

    /// Get object for a handle (returns nullptr if not found).
    SLANG_API ISlangUnknown* getProxy(uint64_t handle) const;

    // ==========================================================================
    // Call Index Access
    // ==========================================================================

    /// Get the number of calls in the loaded index.
    /// Returns 0 if no index is loaded.
    SLANG_API size_t getCallCount() const;

    /// Get an index entry by call number (0-based).
    /// Returns nullptr if index is not loaded or callIndex is out of range.
    SLANG_API const CallIndexEntry* getCallIndexEntry(size_t callIndex) const;

    /// Seek the main stream to a specific call by index.
    /// Returns SLANG_OK on success, SLANG_E_NOT_FOUND if index not loaded,
    /// SLANG_E_INVALID_ARG if callIndex is out of range.
    SLANG_API SlangResult seekToCall(size_t callIndex);

    /// Check if the call index is loaded/available.
    SLANG_API bool hasCallIndex() const { return m_indexStream.getSize() > 0; }

    // ==========================================================================
    // Playback Dispatcher
    // ==========================================================================

    /// Function type for registered playback handlers.
    /// The handler is called in playback mode and should call the appropriate
    /// record() methods to read arguments from the stream, then execute the call.
    using PlaybackHandler = void (*)(ReplayContext& ctx);

    /// Register a playback handler for a function signature.
    /// The signature should match what __FUNCSIG__ or __PRETTY_FUNCTION__ produces.
    SLANG_API void registerHandler(const char* signature, PlaybackHandler handler);

    /// Return the number of registered playback handlers.
    SLANG_API size_t getHandlerCount() const;

    /// Clear the handler dictionary and deallocate its backing storage.
    /// Called from slang_shutdown() to free heap-allocated handler strings
    /// before _CrtDumpMemoryLeaks() runs.
    SLANG_API void resetHandlers();

    /// Re-register the default set of playback handlers.
    /// Call after resetHandlers() to restore replay functionality.
    /// Normal initialization happens automatically via static HandlerRegistrar.
    SLANG_API void registerDefaultHandlers();

    /// Execute the next recorded call from the stream.
    /// Reads the function signature, looks up the handler, and calls it.
    /// Returns true if a call was executed, false if at end of stream.
    SLANG_API bool executeNextCall();

    /// Execute all recorded calls until end of stream.
    SLANG_API void executeAll();

    /// Check if there are more calls to execute.
    SLANG_API bool hasMoreCalls() const { return !m_stream.atEnd(); }

    /// Get the 'this' handle for the current call being executed.
    /// Only valid within a playback handler.
    SLANG_API uint64_t getCurrentThisHandle() const { return m_currentThisHandle; }

    /// Get the 'this' pointer for the current call, cast to the given type.
    /// Only valid within a playback handler.
    /// Note: The handle table stores canonical ISlangUnknown* pointers (via
    /// ProxyBase::toSlangUnknown), which go through TFirstInterface. Since
    /// TFirstInterface is always the first base of the proxy, the pointer
    /// value equals the proxy's address, making this reinterpret_cast safe.
    template<typename T>
    T* getCurrentThis()
    {
        if (m_currentThisHandle == kNullHandle)
            return nullptr;
        auto* unknown = getProxy(m_currentThisHandle);
        return reinterpret_cast<T*>(unknown);
    }

private:
    // Internal recording functions
    SLANG_API void recordRaw(RecordFlag flags, void* data, size_t size);
    SLANG_API void recordTypeId(TypeId id);
    SLANG_API void writeTypeId(TypeId id);
    SLANG_API TypeId readTypeId();
    SLANG_API TypeId readTypeIdFromReference();
    SLANG_API void expectTypeId(TypeId expected);
    SLANG_API void writeIndexEntry();

    // Internal registeration using canonical ISlangUnknown* identities
    SLANG_API uint64_t registerProxyImpl(ISlangUnknown* proxy, ISlangUnknown* implementation);
    SLANG_API void unregisterProxyImpl(ISlangUnknown* proxy);

    /// Remember that the playback dispatcher created `proxy` for an output/return
    /// value and dropped its owning creation reference into a discarded
    /// temporary. Called only during playback; releaseOrphanedPlaybackProxies()
    /// later releases the reference unless the replayed release stream already
    /// destroyed the proxy first (issue #11936).
    SLANG_API void notePlaybackOrphanedProxy(ISlangUnknown* proxy);

    /// Release the still-live orphaned playback proxies recorded by
    /// notePlaybackOrphanedProxy. Proxies whose recorded release stream already
    /// balanced their creation reference destroyed themselves during replay and
    /// were removed by unregisterProxyImpl, so only genuine leaks remain here.
    /// Called at playback teardown (reset/switch/destroy).
    SLANG_API void releaseOrphanedPlaybackProxies();
    SLANG_API ISlangUnknown* getProxyImpl(ISlangUnknown* implementation);
    SLANG_API ISlangUnknown* getImplementationImpl(ISlangUnknown* proxy);
    SLANG_API uint64_t testOnlyRegisterProxyImpl(ISlangUnknown* obj);

    /// Number of orphaned playback references currently tracked for `proxy`, or
    /// 0 if none. Exists so a test can assert the note/unnote balancing
    /// directly: the alternative is inferring it from a sanitizer leak or
    /// double-free report, which only fires on a path some test happens to
    /// drive, and the entry-point wrapping that motivates the `unnote` is not
    /// driven by any of them.
    SLANG_API uint32_t testOnlyGetOrphanedRefCountImpl(ISlangUnknown* proxy) const;
    SLANG_API uint64_t getProxyHandleImpl(ISlangUnknown* obj) const;
    SLANG_API bool isInterfaceRegisteredImpl(ISlangUnknown* obj) const;

    // Initialization and logging
    SLANG_API void ensureInitialized();
    SLANG_API void logCall(const char* signature, void* thisPtr);
    SLANG_API void setupRecordingMirror();
    SLANG_API void closeRecordingMirror();
    SLANG_API static String generateTimestampFolderName();
    SLANG_API void requireReplayArenaAllocation(size_t offset, size_t size);

    /// Record a COM interface pointer (internal implementation).
    template<typename T>
    void recordInterfaceImpl(RecordFlag flags, T*& obj);

    template<typename T, typename CountT>
    T* readArrayInPlayback(RecordFlag flags, CountT& count);

    // Core streams + mutex
    std::recursive_mutex m_mutex;
    ReplayStream m_stream;          ///< Main stream for record/playback
    ReplayStream m_indexStream;     ///< Index stream for call navigation (index.bin)
    ReplayStream m_referenceStream; ///< Reference stream for sync mode comparison
    MemoryArena m_arena;
    size_t m_replayArenaAllocationSize = 0;
    Mode m_mode;
    List<uint8_t> m_compareBuffer; ///< Reusable buffer for sync comparisons

    // Handle tracking: maps objects to handles and back
    Dictionary<ISlangUnknown*, uint64_t> m_objectToHandle;
    Dictionary<uint64_t, ISlangUnknown*> m_handleToObject;
    uint64_t m_nextHandle = kFirstValidHandle;

    // Proxy tracking: maps proxies to implementations and back
    Dictionary<ISlangUnknown*, ISlangUnknown*> m_proxyToImpl;
    Dictionary<ISlangUnknown*, ISlangUnknown*> m_implToProxy;

    // Proxies created by the playback dispatcher for an output/return value.
    // During playback the wrapped object is handed back into a discarded
    // temporary (callWithDefaults passes null outputs), so its owning creation
    // reference has no owner. Keyed by canonical ISlangUnknown* identity -> the
    // number of such orphaned references. releaseOrphanedPlaybackProxies()
    // releases any that the recorded addRef/release stream did not already
    // balance, at playback teardown (issue #11936). Empty outside playback.
    Dictionary<ISlangUnknown*, uint32_t> m_playbackOrphanedProxies;

    // Replay directory management
    String m_replayDirectory = ".slang-replays"; ///< Base directory for replays
    String m_currentReplayPath;                  ///< Current recording session folder

    // TTY logging
    bool m_ttyLogging = false; ///< Whether to log calls to stderr

    // Deferred initialization (to avoid global init order issues with CharEncoding)
    bool m_initialized = false; ///< True after ensureInitialized() has run

    // Map from function signature to handler
    Dictionary<String, PlaybackHandler> m_handlers;

    // Current 'this' handle during playback execution
    uint64_t m_currentThisHandle = kNullHandle;
};

// Template implementations

template<typename T, typename CountT>
T* ReplayContext::readArrayInPlayback(RecordFlag flags, CountT& count)
{
    expectTypeId(TypeId::Array);
    uint64_t arrayCount;
    record(flags, arrayCount);
    size_t countOffset = m_stream.getPosition() - sizeof(arrayCount);
    validateReplayArrayCount(
        m_stream,
        arrayCount,
        uint64_t((std::numeric_limits<CountT>::max)()),
        sizeof(T),
        getReplayArrayMinStreamBytesPerElement<T>(),
        countOffset);
    count = static_cast<CountT>(arrayCount);
    if (arrayCount == 0)
        return nullptr;

    size_t sizeCount = static_cast<size_t>(arrayCount);
    size_t allocationSize = sizeCount * sizeof(T);
    requireReplayArenaAllocation(countOffset, allocationSize);
    T* buf = reinterpret_cast<T*>(allocateReplayArena(allocationSize, alignof(T)));
    for (size_t i = 0; i < sizeCount; ++i)
    {
        new (&buf[i]) T{};
        record(flags, buf[i]);
    }
    return buf;
}

template<typename T, typename CountT>
void ReplayContext::recordArray(RecordFlag flags, T*& arr, CountT& count)
{
    if (m_mode == Mode::Idle)
        return;
    if (isWriting())
    {
        recordTypeId(TypeId::Array);
        uint64_t arrayCount = static_cast<uint64_t>(count);
        record(flags, arrayCount);
        for (uint64_t i = 0; i < arrayCount; ++i)
            record(flags, arr[i]);
    }
    else
    {
        arr = readArrayInPlayback<T>(flags, count);
    }
}

template<typename T, typename CountT>
void ReplayContext::recordArray(RecordFlag flags, const T*& arr, CountT& count)
{
    if (m_mode == Mode::Idle)
        return;
    if (isWriting())
    {
        recordTypeId(TypeId::Array);
        uint64_t arrayCount = static_cast<uint64_t>(count);
        record(flags, arrayCount);
        for (uint64_t i = 0; i < arrayCount; ++i)
            record(flags, const_cast<T&>(arr[i]));
    }
    else
    {
        arr = readArrayInPlayback<T>(flags, count);
    }
}

template<typename EnumT>
void ReplayContext::recordEnum(RecordFlag flags, EnumT& value)
{
    if (m_mode == Mode::Idle)
        return;
    int32_t v = static_cast<int32_t>(value);
    record(flags, v);
    if (isReading())
        value = static_cast<EnumT>(v);
}

template<typename T>
void ReplayContext::recordInterfaceImpl(RecordFlag flags, T*& obj)
{
    if (m_mode == Mode::Idle)
        return;

    bool isInput = hasFlag(flags, RecordFlag::Input) || hasFlag(flags, RecordFlag::ThisPtr);
    bool isOutput = hasFlag(flags, RecordFlag::Output) || hasFlag(flags, RecordFlag::ReturnValue);

    if (isWriting())
    {
        // Recording mode
        if (isInput)
        {
            // An input from the user to a function, should be a proxy they have previously been
            // handed by a wrapping api. Needs unwrapping before returning to hand into main
            // slang api.

            // Handle null
            if (obj == nullptr)
            {
                uint64_t handle = kNullHandle;
                recordHandle(flags, handle);
                return;
            }

            // Normal case: look up handle for tracked object
            uint64_t handle = getProxyHandle(obj);
            recordHandle(flags, handle);

            // Unwrap the proxy to get the underlying implementation
            obj = unwrapObject(obj);
        }
        else if (isOutput)
        {
            // An output from a slang api to be handed back to user. Should be an implementation
            // that needs to be wrapped (or the existing wrapped object needs retrieving) and
            // returned.

            // Wrap the implementation in a proxy
            obj = wrapObject(obj);

            // Output: register object and record handle
            uint64_t handle = getProxyHandle(obj);
            recordHandle(flags, handle);
        }
    }
    else
    {
        // Playback mode
        if (isInput)
        {
            // An input from the user to a function, should be a proxy they have previously been
            // handed by a wrapping api. Needs unwrapping before returning to hand into main
            // slang api.

            // Read handle
            uint64_t handle = kNullHandle;
            recordHandle(flags, handle);

            if (handle == kNullHandle)
            {
                obj = nullptr;
            }
            else
            {
                ISlangUnknown* proxy = getProxy(handle);
                obj = proxy ? toSlangInterface<T>(proxy) : nullptr;
            }

            // Unwrap the proxy to get the underlying implementation
            obj = unwrapObject(obj);
        }
        else if (isOutput)
        {
            // An output from a slang api to be handed back to user. Should be an implementation
            // that needs to be wrapped (or the existing wrapped object needs retrieving) and
            // returned.

            // Wrap the implementation in a proxy
            obj = wrapObject(obj);

            // Output: register object and record handle
            uint64_t handle = getProxyHandle(obj);
            recordHandle(flags, handle);

            // The wrapped proxy's owning reference flows back into a discarded
            // output temporary (callWithDefaults passes null outputs on replay),
            // so nothing releases it. Remember it so playback teardown can
            // release the reference if the recorded release stream did not
            // already balance it (issue #11936).
            notePlaybackOrphanedProxy(toSlangUnknown(obj));
        }
    }
}

} // namespace SlangRecord
