import pytest
from datetime import datetime, timezone
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, empty_checkpoint
from langgraph_checkpoint_timeseries import TimescaleDBSaver, AsyncTimescaleDBSaver

@pytest.fixture
def sync_saver():
    conn_str = "postgresql://postgres:postgres@localhost:5432/postgres"
    with TimescaleDBSaver.from_conn_string(conn_str) as saver:
        saver.setup()
        yield saver

@pytest.fixture
async def async_saver():
    conn_str = "postgresql://postgres:postgres@localhost:5432/postgres"
    async with AsyncTimescaleDBSaver.from_conn_string(conn_str) as saver:
        await saver.setup()
        yield saver

def test_sync_timescaledb_saver(sync_saver: TimescaleDBSaver):
    config = {"configurable": {"thread_id": "test_thread", "checkpoint_ns": ""}}
    checkpoint = empty_checkpoint()
    
    # Test put
    metadata = CheckpointMetadata(source="input", step=0, run_id="test_run", parents={})
    sync_saver.put(config, checkpoint, metadata, {})
    
    # Test get
    retrieved = sync_saver.get_tuple(config)
    assert retrieved is not None
    assert retrieved.checkpoint["id"] == checkpoint["id"]

    # Test put_writes
    sync_saver.put_writes(retrieved.config, [("test_channel", "test_value")], "test_task")
    retrieved_with_writes = sync_saver.get_tuple(retrieved.config)
    
    # Assert writes were retrieved
    assert retrieved_with_writes.pending_writes is not None

@pytest.mark.asyncio
async def test_async_timescaledb_saver():
    conn_str = "postgresql://postgres:postgres@localhost:5432/postgres"
    async with AsyncTimescaleDBSaver.from_conn_string(conn_str) as async_saver:
        await async_saver.setup()
        
        config = {"configurable": {"thread_id": "test_thread_async", "checkpoint_ns": ""}}
        checkpoint = empty_checkpoint()
        
        metadata = CheckpointMetadata(source="input", step=0, run_id="test_run", parents={})
        await async_saver.aput(config, checkpoint, metadata, {})
        
        retrieved = await async_saver.aget_tuple(config)
        assert retrieved is not None
        assert retrieved.checkpoint["id"] == checkpoint["id"]

        await async_saver.aput_writes(retrieved.config, [("test_channel_async", "test_value_async")], "test_task")
        retrieved_with_writes = await async_saver.aget_tuple(retrieved.config)
        
        assert retrieved_with_writes.pending_writes is not None
