import asyncio
import os
import pytest
from stent import Stent

def get_test_backend(test_id: str):
    backend_type = os.environ.get("STENT_TEST_BACKEND", "sqlite")
    
    if backend_type == "postgres":
        # Assumes a running postgres instance, e.g. via docker
        # DSN format: postgres://user:password@host:port/database
        # We append test_id to db name or use schemas? 
        # For simplicity, let's assume a single test db and we clear tables or use unique prefixes if needed.
        # BUT, standard practice for concurrent tests is separate DBs or schemas.
        # Here we will rely on the caller to cleanup or we can use a single DB and init_db ensures tables exist.
        dsn = os.environ.get("STENT_TEST_PG_DSN", "postgres://postgres:postgres@localhost:5432/stent_test")
        backend = Stent.backends.PostgresBackend(dsn)
        return backend
    else:
        db_path = f"test_stent_{test_id}.sqlite"
        return Stent.backends.SQLiteBackend(db_path)

async def clear_test_backend(backend):
    """Clears tables but keeps connection open if possible"""
    if hasattr(backend, "pool") and backend.pool:
         async with backend.pool.acquire() as conn:
             # Check if tables exist before truncating
              exists = await conn.fetchval(
                 "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'executions')"
              )
              if exists:
                 await conn.execute("TRUNCATE TABLE executions, execution_progress, tasks, dead_tasks, cache, idempotency, execution_counters, execution_state CASCADE")

async def cleanup_test_backend(backend):
    # First close the backend connection
    if hasattr(backend, "close"):
        await backend.close()
    
    if hasattr(backend, "db_path"):
        if not os.path.exists(backend.db_path):
            return

        for attempt in range(20):
            try:
                os.remove(backend.db_path)
                return
            except PermissionError:
                # Windows keeps SQLite files locked briefly after connections close.
                if attempt == 19:
                    return
                await asyncio.sleep(0.05 * (attempt + 1))
    elif hasattr(backend, "pool") and backend.pool:
         # For postgres, we might want to truncate tables
         async with backend.pool.acquire() as conn:
             await conn.execute("TRUNCATE TABLE executions, execution_progress, tasks, dead_tasks, cache, idempotency, execution_counters, execution_state CASCADE")
         await backend.pool.close()
