import time
import functools
import weakref
import threading
from typing import Any

import google.cloud.firestore_v1

WORKAROUND = False

def test_watch_cleanup(client: google.cloud.firestore_v1.Client) -> None:
    """Test Watch object garbage collection"""

    # Create a ref to any document (doesn't need to exist).
    doc_ref = client.collection('some_collection').document('some_doc_id')

    # Start listening for changes.
    doc_watch = doc_ref.on_snapshot(on_snapshot)

    # Give it a moment.
    time.sleep(3.0)

    if WORKAROUND:
        # Need to grab these here before unsubscribe clears them.
        rpc = doc_watch._rpc
        consumer = doc_watch._consumer

    # Stop listening for changes.
    doc_watch.unsubscribe()

    if WORKAROUND:
        # From looking at which references were keeping doc_watch alive,
        # I found that clearing these fields after the unsubscribe allows
        # it to go down cleanly.
        rpc._initial_request = None
        rpc._callbacks = []
        consumer._on_response = None
        doc_watch._snapshot_callback = None

    # Lastly, kick off a thread to keep an eye on doc_watch until it dies.
    threading.Thread(
        target=functools.partial(watch_the_watcher, weakref.ref(doc_watch))
    ).start()

def on_snapshot(*args: Any) -> None:
    """Listener callback"""
    print('got snapshot')

def watch_the_watcher(doc_watch_weak_ref: weakref.ReferenceType) -> None:
    """Spin and wait for doc_watch to die."""

    starttime = time.time()
    while doc_watch_weak_ref() is not None:
        print(f'doc_watch is still alive ({time.time()-starttime:.0f}s)')
        time.sleep(5.0)
    print('doc_watch IS DEAD!!!')