from __future__ import annotations

import pytest
import pytest_asyncio

from db import Database, UserAlreadyQueued


@pytest_asyncio.fixture
async def db(tmp_path):
    database = Database(tmp_path / "bot.db")
    await database.connect()
    yield database
    await database.close()


async def test_create_and_read(db: Database) -> None:
    request = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="apple",
        input_file_path=None,
        seed=42,
    )
    assert request.id > 0
    assert request.status == "queued"
    assert request.mode == "t2i"
    assert request.seed == 42

    loaded = await db.get_request(request.id)
    assert loaded == request


async def test_one_active_request_per_user(db: Database) -> None:
    first = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="first",
        input_file_path=None,
        seed=1,
    )
    with pytest.raises(UserAlreadyQueued):
        await db.create_request(
            user_id=1,
            mode="t2i",
            prompt_text="second",
            input_file_path=None,
            seed=2,
        )

    await db.set_done(first.id, "/tmp/out.png")
    second = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="second",
        input_file_path=None,
        seed=3,
    )
    assert second.id != first.id


async def test_positions(db: Database) -> None:
    first = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="first",
        input_file_path=None,
        seed=1,
    )
    second = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="second",
        input_file_path=None,
        seed=2,
    )
    assert await db.count_ahead(first.id) == 0
    assert await db.count_ahead(second.id) == 1

    await db.set_running(first.id, "prompt-1")
    assert await db.count_ahead(second.id) == 1

    await db.set_done(first.id, "/tmp/out.png")
    assert await db.count_ahead(second.id) == 0


async def test_recover_marks_running_failed_and_returns_queued(db: Database) -> None:
    running = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="running",
        input_file_path=None,
        seed=1,
    )
    queued = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="queued",
        input_file_path=None,
        seed=2,
    )
    await db.set_running(running.id, "prompt-1")

    recovered = await db.recover()

    assert [item.id for item in recovered] == [queued.id]
    failed = await db.get_request(running.id)
    assert failed is not None
    assert failed.status == "failed"
    assert failed.finished_at is not None