# test_pickledb.py
import os
import json
import asyncio
from pathlib import Path

import pytest

from pickledb import PickleDB, PickleDBSQLite  # adjust this import if your module name is different


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def make_tmp_path(tmp_path, name="db.json") -> Path:
    """Helper to get a db path inside pytest's tmp_path fixture."""
    return tmp_path / name


# ---------------------------------------------------------------------------
# Basic sync usage (PickleDB)
# ---------------------------------------------------------------------------

def test_sync_set_get_save_and_reload(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    # sync .set() & .save() use dualmethod (internally asyncio.run)
    assert db.set("foo", {"bar": 1}) is True
    assert db.set("num", 123) is True
    assert db.save() is True

    # New instance loads from same file
    db2 = PickleDB(str(db_path))
    db2.load()
    assert db2.get("foo") == {"bar": 1}
    assert db2.get("num") == 123


def test_sync_get_default_and_missing_key(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    db.set("exists", "yep")
    db.save()

    db2 = PickleDB(str(db_path))
    db2.load()

    assert db2.get("exists", default="nope") == "yep"
    assert db2.get("missing", default="nope") == "nope"
    # default=None is implicit
    assert db2.get("also_missing") is None


def test_sync_remove_and_all(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    db.set("a", 1)
    db.set("b", 2)
    db.set("c", 3)

    keys = db.all()
    assert set(keys) == {"a", "b", "c"}

    # remove existing key returns True
    assert db.remove("b") is True
    # removing again returns False
    assert db.remove("b") is False

    keys_after = db.all()
    assert set(keys_after) == {"a", "c"}


def test_sync_purge_clears_database(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    db.set("x", 1)
    db.set("y", 2)
    assert set(db.all()) == {"x", "y"}

    assert db.purge() is True
    assert db.all() == []

    db.save()

    # After reload, still empty
    db2 = PickleDB(str(db_path))
    db2.load()
    assert db2.all() == []


# ---------------------------------------------------------------------------
# Basic async usage (PickleDB)
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_async_set_get_save_and_reload(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    await db.set("foo", {"bar": 1})
    await db.set("num", 123)
    await db.save()

    db2 = PickleDB(str(db_path))
    await db2.load()
    assert await db2.get("foo") == {"bar": 1}
    assert await db2.get("num") == 123


@pytest.mark.asyncio
async def test_async_get_default(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    await db.set("exists", "yep")
    await db.save()

    db2 = PickleDB(str(db_path))
    await db2.load()

    assert await db2.get("exists", default="nope") == "yep"
    assert await db2.get("missing", default="nope") == "nope"
    assert await db2.get("also_missing") is None


@pytest.mark.asyncio
async def test_async_remove_all_purge(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    await db.set("a", 1)
    await db.set("b", 2)
    await db.set("c", 3)

    keys = await db.all()
    assert set(keys) == {"a", "b", "c"}

    assert await db.remove("b") is True
    assert await db.remove("b") is False
    keys_after = await db.all()
    assert set(keys_after) == {"a", "c"}

    assert await db.purge() is True
    assert await db.all() == []


# ---------------------------------------------------------------------------
# Context managers (PickleDB)
# ---------------------------------------------------------------------------

def test_sync_context_manager_saves_on_success(tmp_path):
    db_path = make_tmp_path(tmp_path)

    with PickleDB(str(db_path)) as db:
        db.set("inside", "context")
        # __exit__ should call .save() because no exception

    # New instance: must see saved data
    db2 = PickleDB(str(db_path))
    db2.load()
    assert db2.get("inside") == "context"


def test_sync_context_manager_does_not_save_on_exception(tmp_path):
    db_path = make_tmp_path(tmp_path)

    try:
        with PickleDB(str(db_path)) as db:
            db.set("inside", "context")
            raise RuntimeError("boom")
    except RuntimeError:
        pass

    # Because __exit__ sees an exception, it should NOT call .save()
    # Therefore file may not exist, or be empty.
    if db_path.exists():
        content = db_path.read_bytes()
        assert content == b""  # empty file
    else:
        # No file is also acceptable
        assert not db_path.exists()


@pytest.mark.asyncio
async def test_async_context_manager_saves_on_success(tmp_path):
    db_path = make_tmp_path(tmp_path)

    async with PickleDB(str(db_path)) as db:
        await db.set("inside", "async_context")

    db2 = PickleDB(str(db_path))
    await db2.load()
    assert await db2.get("inside") == "async_context"


@pytest.mark.asyncio
async def test_async_context_manager_does_not_save_on_exception(tmp_path):
    db_path = make_tmp_path(tmp_path)

    with pytest.raises(RuntimeError):
        async with PickleDB(str(db_path)) as db:
            await db.set("inside", "async_context")
            raise RuntimeError("boom")

    if db_path.exists():
        content = db_path.read_bytes()
        assert content == b""
    else:
        assert not db_path.exists()


# ---------------------------------------------------------------------------
# Edge cases: missing file, empty file, atomic save (PickleDB)
# ---------------------------------------------------------------------------

def test_load_on_missing_file_returns_empty_db(tmp_path):
    db_path = make_tmp_path(tmp_path)
    assert not db_path.exists()

    db = PickleDB(str(db_path))
    db.load()
    assert db.all() == []
    # After save, file should exist with {} or similar JSON
    db.save()
    assert db_path.exists()

    raw = db_path.read_bytes()
    data = json.loads(raw.decode("utf-8")) if raw else {}
    assert data == {}


def test_load_on_empty_file_returns_empty_db(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db_path.write_text("")  # existing but empty file

    db = PickleDB(str(db_path))
    # load() treats 0-byte file as empty db
    db.load()
    assert db.all() == []


def test_save_uses_temp_file_and_is_atomic(tmp_path):
    db_path = make_tmp_path(tmp_path)
    tmp_name = str(db_path) + ".tmp"

    db = PickleDB(str(db_path))
    db.set("a", 1)
    db.save()

    # temp file should not be left behind
    assert not os.path.exists(tmp_name)

    # file should be valid JSON (orjson output is still standard JSON)
    raw = db_path.read_bytes()
    data = json.loads(raw.decode("utf-8"))
    assert data == {"a": 1}


# ---------------------------------------------------------------------------
# Concurrency tests (PickleDB)
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_async_concurrent_sets_are_all_present(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    async def writer(start, end):
        for i in range(start, end):
            await db.set(f"key_{i}", i)

    # Kick off a few concurrent writers
    tasks = [
        asyncio.create_task(writer(0, 300)),
        asyncio.create_task(writer(300, 600)),
        asyncio.create_task(writer(600, 1000)),
    ]
    await asyncio.gather(*tasks)

    keys = await db.all()
    assert len(keys) == 1000
    # spot check
    assert await db.get("key_0") == 0
    assert await db.get("key_999") == 999


@pytest.mark.asyncio
async def test_async_concurrent_gets_and_sets(tmp_path):
    db_path = make_tmp_path(tmp_path)
    db = PickleDB(str(db_path))

    await db.set("counter", 0)

    async def incrementer(n_times):
        for _ in range(n_times):
            # Not atomic on purpose, but we are testing that we don't crash
            current = await db.get("counter", 0)
            await db.set("counter", current + 1)

    tasks = [asyncio.create_task(incrementer(100)) for _ in range(10)]
    await asyncio.gather(*tasks)

    # We *expect* the final value to be <= 1000 because increments are not atomic.
    # The important part is no exceptions and a sensible final value.
    final_value = await db.get("counter")
    assert 0 < final_value <= 1000


# ---------------------------------------------------------------------------
# Stress test: 1,000,000 key-value pairs (PickleDB)
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
@pytest.mark.stress
async def test_stress_one_million_entries(tmp_path):
    """
    Stress test inserting and retrieving 1,000,000 key-value pairs.

    NOTE:
    - This test is intentionally heavy and may take a long time.
    - Run it explicitly: `pytest -m stress`.
    """
    N = 1_000_000
    db_path = make_tmp_path(tmp_path, "stress_db.json")
    db = PickleDB(str(db_path))

    # Insert 1M entries
    for i in range(N):
        await db.set(f"key_{i}", i)

    # Retrieve 1M entries
    for i in range(N):
        value = await db.get(f"key_{i}")
        assert value == i

    await db.save()
    assert db_path.exists()


# ---------------------------------------------------------------------------
# Basic tests for PickleDBSQLite
# ---------------------------------------------------------------------------

def test_sqlite_sync_set_get_and_all(tmp_path):
    sqlite_path = tmp_path / "kv.sqlite3"
    kv = PickleDBSQLite(str(sqlite_path))

    key1 = kv.set(None, {"foo": "bar"})
    key2 = kv.set("explicit", [1, 2, 3])

    assert isinstance(key1, str)
    assert key2 == "explicit"

    assert kv.get(key1) == {"foo": "bar"}
    assert kv.get("explicit") == [1, 2, 3]

    keys = set(kv.all())
    assert key1 in keys
    assert "explicit" in keys

    kv.close()


@pytest.mark.asyncio
async def test_sqlite_async_set_get_and_purge(tmp_path):
    sqlite_path = tmp_path / "kv_async.sqlite3"
    kv = PickleDBSQLite(str(sqlite_path))

    key = await kv.set(None, {"async": True})
    assert await kv.get(key) == {"async": True}

    await kv.set("x", 1)
    await kv.set("y", 2)
    keys = set(await kv.all())
    assert key in keys
    assert {"x", "y"} <= keys

    assert await kv.remove("x") is True
    assert await kv.remove("x") is False

    assert await kv.purge() is True
    assert await kv.all() == []

    await kv.close()