import asyncio
import os
import sys
from pathlib import Path
import aiosqlite

DB_FILE = Path(__file__).resolve().parent.parent / "test_sessions.db"

async def init_test_db(db_path, use_wal=False):
    if db_path.exists():
        try:
            os.remove(db_path)
            for suffix in ["-wal", "-journal", "-shm"]:
                p = Path(str(db_path) + suffix)
                if p.exists():
                    os.remove(p)
        except OSError:
            pass

    async with aiosqlite.connect(db_path) as db:
        if use_wal:
            await db.execute("PRAGMA journal_mode=WAL")
            await db.execute("PRAGMA busy_timeout=10000")
            await db.execute("PRAGMA synchronous=NORMAL")
        else:
            await db.execute("PRAGMA journal_mode=DELETE")
            await db.execute("PRAGMA busy_timeout=0")
        
        await db.execute("""
            CREATE TABLE IF NOT EXISTS chat_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                role TEXT NOT NULL,
                text TEXT NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        """)
        await db.commit()

async def writer_task(db_path, worker_id, use_wal, sleep_time=0.1):
    timeout = 10.0 if use_wal else 0.0
    async with aiosqlite.connect(db_path, timeout=timeout) as db:
        if not use_wal:
            await db.execute("PRAGMA busy_timeout=0")
        else:
            await db.execute("PRAGMA busy_timeout=10000")
            
        await db.execute("BEGIN IMMEDIATE")
        await db.execute(
            "INSERT INTO chat_history (session_id, role, text) VALUES (?, ?, ?)",
            (f"session_{worker_id}", "user", f"message {worker_id}")
        )
        await asyncio.sleep(sleep_time)
        await db.commit()

async def run_concurrency_test(use_wal):
    await init_test_db(DB_FILE, use_wal=use_wal)
    
    tasks = [writer_task(DB_FILE, i, use_wal) for i in range(5)]
    results = await asyncio.gather(*tasks, return_exceptions=True)
    
    errors = [r for r in results if isinstance(r, Exception)]
    
    if DB_FILE.exists():
        try:
            os.remove(DB_FILE)
            for suffix in ["-wal", "-journal", "-shm"]:
                p = Path(str(DB_FILE) + suffix)
                if p.exists():
                    os.remove(p)
        except OSError:
            pass
            
    return errors

async def main():
    print("--- Testing CONCURRENCY WITHOUT WAL (Should FAIL/RED) ---")
    errors = await run_concurrency_test(use_wal=False)
    if not errors:
        print("RED Test failed: Expected database locks but none occurred.")
        sys.exit(1)
    else:
        print(f"RED Test PASSED: Successfully caught {len(errors)} lock exceptions as expected!")
        for e in errors[:2]:
            print(f"Sample error: {type(e).__name__}: {e}")
            
    print("\n--- Testing CONCURRENCY WITH WAL (Should PASS/GREEN) ---")
    errors = await run_concurrency_test(use_wal=True)
    if errors:
        print(f"GREEN Test failed: Got lock exceptions under WAL mode! {errors}")
        sys.exit(1)
    else:
        print("GREEN Test PASSED: SQLite concurrent writes completed successfully under WAL!")
        sys.exit(0)

if __name__ == "__main__":
    asyncio.run(main())