"""
Ch11 配套代码 3 / 4 —— 看门狗续期机制

业务执行 6 秒，但锁 TTL 只有 2 秒。
对比两种实现：
  ① 朴素锁：业务跑完锁早就被别人抢走
  ② 看门狗锁：守护线程每 0.6s 自动续期，业务安全完成

看门狗实现要点：
  - 后台 daemon 线程，定期 PEXPIRE
  - 续期前用 Lua 校验持有者
  - 主线程 release 时 Event.set() 通知看门狗退出
"""

import time
import uuid
import threading

try:
    import redis
except ImportError:
    print("请先 pip install redis"); raise SystemExit(1)


RELEASE_LUA = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('DEL', KEYS[1])
else return 0 end
"""

RENEW_LUA = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('PEXPIRE', KEYS[1], ARGV[2])
else return 0 end
"""


def now() -> str:
    return time.strftime('%H:%M:%S', time.localtime()) + f".{int(time.time()*1000)%1000:03d}"


class WatchdogLock:
    """ 带看门狗的安全锁。"""

    def __init__(self, client: "redis.Redis", key: str, ttl_ms: int = 2000):
        self.client = client
        self.key = key
        self.ttl_ms = ttl_ms
        self.token = uuid.uuid4().hex
        self._stop = None
        self._thread = None

    def acquire(self, watchdog: bool = True) -> bool:
        ok = self.client.set(self.key, self.token, nx=True, px=self.ttl_ms)
        if not ok:
            return False
        if watchdog:
            self._start_watchdog()
        return True

    def release(self) -> None:
        self._stop_watchdog()
        self.client.eval(RELEASE_LUA, 1, self.key, self.token)

    def _start_watchdog(self) -> None:
        self._stop = threading.Event()
        interval = self.ttl_ms / 3 / 1000.0
        client = self.client
        key, token, ttl = self.key, self.token, self.ttl_ms

        def loop():
            while not self._stop.wait(interval):
                ok = client.eval(RENEW_LUA, 1, key, token, ttl)
                if ok == 1:
                    print(f"  [{now()}] 🐶 看门狗续期 PEXPIRE {ttl}ms ✓")
                else:
                    print(f"  [{now()}] 🐶 锁已不在自己手上，看门狗退出")
                    return

        self._thread = threading.Thread(target=loop, daemon=True)
        self._thread.start()

    def _stop_watchdog(self) -> None:
        if self._stop is not None:
            self._stop.set()
            if self._thread is not None:
                self._thread.join(timeout=1)


# -------------------------------------------------------------
# 演示一：无看门狗 → 业务超时被抢
# -------------------------------------------------------------
def demo_no_watchdog() -> None:
    print("\n" + "=" * 60)
    print("演示 ①：无看门狗（业务 6s × TTL 2s → 中途被抢）")
    print("=" * 60)
    r = redis.Redis(host="127.0.0.1", port=6379, decode_responses=True)
    key = "demo:wd:no"
    r.delete(key)

    a = WatchdogLock(r, key, ttl_ms=2000)
    assert a.acquire(watchdog=False)
    print(f"  [{now()}] [A] 拿到锁，TTL=2s，业务开始（要跑 6s）")

    def thief():
        time.sleep(2.5)  # 等 A 的锁过期
        b = WatchdogLock(r, key, ttl_ms=2000)
        if b.acquire(watchdog=False):
            print(f"  [{now()}] [B] ⚠️ 趁 A 锁过期，把锁抢走了！")
            time.sleep(0.5)
            b.release()

    t = threading.Thread(target=thief)
    t.start()

    for i in range(6):
        time.sleep(1)
        owner = r.get(key)
        flag = "✓ 还是 A" if owner == a.token else ("✗ 已被 B 抢" if owner else "已过期")
        print(f"  [{now()}] [A] 业务进行中 {i+1}/6s … 锁所有者: {flag}")
    t.join()
    a.release()
    print(f"  [{now()}] [A] 业务完成 ❌ —— 中途锁被夺走，临界区被破坏")
    r.delete(key)


# -------------------------------------------------------------
# 演示二：有看门狗 → 安全完成
# -------------------------------------------------------------
def demo_with_watchdog() -> None:
    print("\n" + "=" * 60)
    print("演示 ②：有看门狗（业务 6s × TTL 2s × 续期间隔 ~0.67s）")
    print("=" * 60)
    r = redis.Redis(host="127.0.0.1", port=6379, decode_responses=True)
    key = "demo:wd:yes"
    r.delete(key)

    a = WatchdogLock(r, key, ttl_ms=2000)
    assert a.acquire(watchdog=True)
    print(f"  [{now()}] [A] 拿到锁，TTL=2s，业务开始（要跑 6s）")

    def thief():
        for i in range(6):
            time.sleep(1)
            b = WatchdogLock(r, key, ttl_ms=2000)
            got = b.acquire(watchdog=False)
            if got:
                print(f"  [{now()}] [B] ⚠️ 抢锁成功（不该发生！）")
                b.release()

    t = threading.Thread(target=thief)
    t.start()

    for i in range(6):
        time.sleep(1)
        owner = r.get(key)
        ok = "✓ A 持锁" if owner == a.token else "✗ 锁丢了"
        print(f"  [{now()}] [A] 业务 {i+1}/6s … {ok}")

    t.join()
    a.release()
    print(f"  [{now()}] [A] 业务完成 ✅ —— 看门狗保住了临界区")
    r.delete(key)


if __name__ == "__main__":
    try:
        redis.Redis(host="127.0.0.1", port=6379).ping()
    except redis.ConnectionError as e:
        print(f"❌ Redis 连接失败：{e}")
        raise SystemExit(1)

    demo_no_watchdog()
    demo_with_watchdog()
