#include "CcosLock.h"
#include <pthread.h>
#include <unistd.h>
#include <map>
#include <atomic>
#include <cerrno>  // 包含 ETIMEDOUT 定义

// 进程锁实现
class ProcessLock {
    pthread_mutex_t m_Mutex = PTHREAD_MUTEX_INITIALIZER;
    pthread_cond_t m_Cond = PTHREAD_COND_INITIALIZER;
    bool m_Locked = false;
    
public:
    DWORD Lock(DWORD timeout) {
        pthread_mutex_lock(&m_Mutex);
        
        if (!m_Locked) {
            m_Locked = true;
            pthread_mutex_unlock(&m_Mutex);
            return WAIT_OBJECT_0;
        }
        
        if (timeout == 0) {
            pthread_mutex_unlock(&m_Mutex);
            return WAIT_TIMEOUT;
        }
        
        struct timespec ts;
        if (timeout != INFINITE) {
            clock_gettime(CLOCK_REALTIME, &ts);
            ts.tv_sec += timeout / 1000;
            ts.tv_nsec += (timeout % 1000) * 1000000;
            if (ts.tv_nsec >= 1000000000) {
                ts.tv_sec++;
                ts.tv_nsec -= 1000000000;
            }
        }
        
        while (m_Locked) {
            int ret = (timeout == INFINITE) 
                ? pthread_cond_wait(&m_Cond, &m_Mutex)
                : pthread_cond_timedwait(&m_Cond, &m_Mutex, &ts);
            
            if (ret == ETIMEDOUT) {  // 现在 ETIMEDOUT 已定义
                pthread_mutex_unlock(&m_Mutex);
                return WAIT_TIMEOUT;
            }
        }
        
        m_Locked = true;
        pthread_mutex_unlock(&m_Mutex);
        return WAIT_OBJECT_0;
    }

    void Unlock() {
        pthread_mutex_lock(&m_Mutex);
        m_Locked = false;
        pthread_cond_signal(&m_Cond);
        pthread_mutex_unlock(&m_Mutex);
    }
};

static ProcessLock g_ProcessLock;

// 全局数据结构
static pthread_mutex_t g_MapMutex = PTHREAD_MUTEX_INITIALIZER;
static std::map<UINT64, DWORD> g_LockedMap;        // 对象 -> 线程ID
static std::map<UINT64, CcosLock*> g_ObjectMap;    // 对象指针映射

// 注册/注销锁对象
void RegisterLock(CcosLock* p) {
    pthread_mutex_lock(&g_MapMutex);
    g_ObjectMap[reinterpret_cast<UINT64>(p)] = p;
    pthread_mutex_unlock(&g_MapMutex);
}

void UnregisterLock(CcosLock* p) {
    pthread_mutex_lock(&g_MapMutex);
    UINT64 key = reinterpret_cast<UINT64>(p);
    g_ObjectMap.erase(key);
    g_LockedMap.erase(key);
    pthread_mutex_unlock(&g_MapMutex);
}

// CcosLock 成员函数实现
CcosLock::CcosLock() {
    pthread_mutex_init(&m_Mutex, nullptr);
    pthread_cond_init(&m_Cond, nullptr);
    RegisterLock(this);
}

CcosLock::~CcosLock() {
    UnregisterLock(this);
    pthread_mutex_destroy(&m_Mutex);
    pthread_cond_destroy(&m_Cond);
}

bool CcosLock::CcosInternal_EnterCriticalSection(DWORD TryCount) {
    const DWORD tid = GetCurrentThreadId();
    DWORD expected = 0;
    
    // 尝试获取锁
    if (__atomic_compare_exchange_n(&m_InterLock, &expected, tid, 
                                   false, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED)) {
        m_RefCount = 1;
        pthread_mutex_lock(&g_MapMutex);
        g_LockedMap[reinterpret_cast<UINT64>(this)] = tid;
        pthread_mutex_unlock(&g_MapMutex);
        return true;
    }
    
    // 检查是否已持有锁
    if (expected == tid) {
        m_RefCount++;
        return true;
    }
    
    // 尝试多次获取锁
    while (TryCount-- > 0) {
        expected = 0;
        if (__atomic_compare_exchange_n(&m_InterLock, &expected, tid, 
                                       false, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED)) {
            m_RefCount = 1;
            pthread_mutex_lock(&g_MapMutex);
            g_LockedMap[reinterpret_cast<UINT64>(this)] = tid;
            pthread_mutex_unlock(&g_MapMutex);
            return true;
        }
        usleep(1000); // 1ms 延迟
    }
    return false;
}

void CcosLock::CcosInternal_LeaveCriticalSection() {
    const DWORD tid = GetCurrentThreadId();
    if (m_InterLock == tid) {
        if (--m_RefCount == 0) {
            pthread_mutex_lock(&g_MapMutex);
            g_LockedMap.erase(reinterpret_cast<UINT64>(this));
            pthread_mutex_unlock(&g_MapMutex);
            __atomic_store_n(&m_InterLock, 0, __ATOMIC_RELEASE);
        }
    }
}

DWORD CcosLock::Thread_Lock(DWORD timeout) {
    if (CcosInternal_EnterCriticalSection(1)) {
        return WAIT_OBJECT_0;
    }

    const DWORD start = GetTickCount();
    while (true) {
        const DWORD elapsed = GetTickCount() - start;
        if (timeout != INFINITE && elapsed >= timeout) {
            return WAIT_TIMEOUT;
        }
        
        const DWORD remaining = (timeout == INFINITE) ? 100 : (timeout - elapsed);
        const DWORD waitTime = std::min<DWORD>(remaining, 100);
        const DWORD ret = Thread_WaitUnlockNotify(waitTime);
        
        if (ret == WAIT_OBJECT_0 && CcosInternal_EnterCriticalSection(1)) {
            return WAIT_OBJECT_0;
        }
    }
}

void CcosLock::Thread_UnLock() {
    CcosInternal_LeaveCriticalSection();
}

bool CcosLock::Thread_Clear(DWORD tid) {
    if (m_InterLock == tid) {
        m_RefCount = 0;
        pthread_mutex_lock(&g_MapMutex);
        g_LockedMap.erase(reinterpret_cast<UINT64>(this));
        pthread_mutex_unlock(&g_MapMutex);
        __atomic_store_n(&m_InterLock, 0, __ATOMIC_RELEASE);
        pthread_cond_broadcast(&m_Cond);
        return true;
    }
    return false;
}

DWORD CcosLock::Thread_WaitUnlockNotify(DWORD timeout) {
    pthread_mutex_lock(&m_Mutex);
    
    struct timespec ts;
    if (timeout != INFINITE) {
        clock_gettime(CLOCK_REALTIME, &ts);
        ts.tv_sec += timeout / 1000;
        ts.tv_nsec += (timeout % 1000) * 1000000;
        if (ts.tv_nsec >= 1000000000) {
            ts.tv_sec++;
            ts.tv_nsec -= 1000000000;
        }
    }
    
    while (m_InterLock != 0) {
        int ret = (timeout == INFINITE)
            ? pthread_cond_wait(&m_Cond, &m_Mutex)
            : pthread_cond_timedwait(&m_Cond, &m_Mutex, &ts);
        
        if (ret == ETIMEDOUT) break;  // 现在 ETIMEDOUT 已定义
    }
    
    pthread_mutex_unlock(&m_Mutex);
    return (m_InterLock == 0) ? WAIT_OBJECT_0 : WAIT_TIMEOUT;
}

int CcosLock::FindProcessThreads(DWORD tid) {
    return IsThreadAlive(tid);
}

bool CcosLock::Try_Clear_DeadThread() {
    return (m_InterLock != 0 && !FindProcessThreads(m_InterLock))
        ? Thread_Clear(m_InterLock) : false;
}

// 全局清理函数
void CleanupForThread(DWORD tid) {
    pthread_mutex_lock(&g_MapMutex);
    
    auto it = g_LockedMap.begin();
    while (it != g_LockedMap.end()) {
        if (it->second == tid) {
            if (g_ObjectMap.find(it->first) != g_ObjectMap.end()) {
                g_ObjectMap[it->first]->Thread_Clear(tid);
            }
            it = g_LockedMap.erase(it);
        } else {
            ++it;
        }
    }
    
    pthread_mutex_unlock(&g_MapMutex);
}

// C接口函数实现
bool Ccos_ThreadLock(volatile UINT64 *pLock) {
    volatile DWORD* lockPtr = reinterpret_cast<volatile DWORD*>(pLock);
    const DWORD tid = GetCurrentThreadId();
    DWORD expected = 0;
    
    if (__atomic_compare_exchange_n(lockPtr, &expected, tid, 
                                   false, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED)) {
        lockPtr[1] = 1; // refcount
        return true;
    }
    
    if (expected == tid) {
        lockPtr[1]++; // 增加引用计数
        return true;
    }
    
    return false;
}

void Ccos_ThreadUnLock(volatile UINT64 *pLock) {
    volatile DWORD* lockPtr = reinterpret_cast<volatile DWORD*>(pLock);
    const DWORD tid = GetCurrentThreadId();
    
    if (lockPtr[0] == tid) {
        if (--lockPtr[1] == 0) { // 减少引用计数
            __atomic_store_n(lockPtr, 0, __ATOMIC_RELEASE);
        }
    }
}

void Ccos_ThreadClearForTid(volatile UINT64 *pLock, DWORD tid) {
    volatile DWORD* lockPtr = reinterpret_cast<volatile DWORD*>(pLock);
    
    if (lockPtr[0] == tid) {
        lockPtr[1] = 0; // 重置引用计数
        __atomic_store_n(lockPtr, 0, __ATOMIC_RELEASE);
    }
}

DWORD Proc_Lock(DWORD timeout) {
    return g_ProcessLock.Lock(timeout);
}

void Proc_UnLock() {
    g_ProcessLock.Unlock();
}

DWORD Thread_GetUniqTick() {
    static std::atomic<DWORD> counter(1);
    return counter.fetch_add(1, std::memory_order_relaxed);
}