#!/usr/bin/env python3

import sys
import os
import asyncio
import tempfile
import aiosqlite
import zstandard
from pypaste.server import Paste, Storage, keygen
from pypaste.server.sqlite import Sqlite
from pypaste.server.s3 import S3
from datetime import datetime
from pathlib import Path


def truncate(path: Path) -> None:
    with open(path, "w") as f:
        f.truncate(0)


async def test_storage_use(storage: Storage) -> None:
    content = "hello pypaste"
    length = len(zstandard.compress(content.encode()))

    for x in range(5):
        paste = Paste(datetime.now(), None, content)
        key = keygen(6)
        await storage.insert(paste, key)

    assert (use := await storage.storage_use()) is not None and use == length * 5


async def test_oldest(storage: Storage) -> None:
    pastes = []
    for x in range(5):
        paste = Paste(datetime.now(), None, "hello pypaste")
        key = keygen(6)
        pastes.append(key)
        await storage.insert(paste, key)

    assert (oldest := await storage.oldest()) is not None and oldest == pastes[0]


async def test_exists_but_not_in_our_table(storage: Storage) -> None:
    key = keygen(6)

    await storage.connection.execute(
        "insert into pastes values(?, ?, ?, ?, ?)",
        (key.data, key.length, datetime.now().isoformat(), None, bytes()),
    )

    assert not await storage.exists(key)


async def test_exists(storage: Storage) -> None:
    dt = datetime.now()
    key = keygen(6)

    await storage.insert(Paste(dt, "test", "hello world"), key)

    assert await storage.exists(key)

    assert not await storage.exists(keygen(6))


async def test_delete(storage: Storage) -> None:
    dt = datetime.now()
    key = keygen(6)

    await storage.insert(Paste(dt, "test", "hello world"), key)

    assert await storage.exists(key)

    await storage.delete(key)

    assert not await storage.exists(key)


async def test_insert_retrieve(storage: Storage) -> None:
    dt = datetime.now()
    key = keygen(6)

    await storage.insert(Paste(dt, "test", "hello world"), key)

    paste = await storage.retrieve(key)

    assert paste is not None
    assert paste.dt == dt
    assert paste.syntax == "test"
    assert paste.text == "hello world"


async def init_db(connection: aiosqlite.Connection) -> None:
    await connection.execute(
        (
            "create table pastes("
            "key blob,"
            "key_length int,"
            "datetime text,"
            "size int,"
            "syntax text"
            ")"
        )
    )


async def test_sqlite(tests):
    for test in tests:
        with tempfile.TemporaryDirectory() as tmpdir:
            database = Path(tmpdir) / "pastes.sqlite"
            truncate(database)

            async with aiosqlite.connect(database) as connection:
                await init_db(connection)

                sqlite = Sqlite(connection)

                await sqlite.setup()

                await test(sqlite)


async def test_s3(tests):
    for test in tests:
        with tempfile.TemporaryDirectory() as tmpdir:
            database = Path(tmpdir) / "pastes.sqlite"
            truncate(database)

            async with aiosqlite.connect(database) as connection:
                await init_db(connection)

                s3 = S3(
                    connection,
                    os.environ["PYPASTE_TEST_ENDPOINT"],
                    os.environ["PYPASTE_TEST_REGION"],
                    os.environ["PYPASTE_TEST_BUCKET"],
                    os.environ["PYPASTE_TEST_ACCESS_KEY"],
                    os.environ["PYPASTE_TEST_SECRET_KEY"],
                )

                await s3.setup()

                await test(s3)


async def main() -> int:
    tests = [
        test_insert_retrieve,
        test_delete,
        test_exists,
        test_exists_but_not_in_our_table,
        test_oldest,
        test_storage_use,
    ]

    await test_sqlite(tests)

    try:
        os.environ["PYPASTE_TEST_S3"]
        await test_s3(tests)
    except KeyError:
        pass

    return 0


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