from __future__ import annotations

from datetime import timedelta

import pytest

from db import Database, UserAlreadyQueued, msk_today


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


async def test_history_filters_and_stats(db: Database) -> None:
    first = await db.create_request(
        user_id=1,
        username="alice",
        mode="t2i",
        prompt_text="apple",
        input_file_path=None,
        seed=1,
    )
    await db.set_done(first.id, "/tmp/1.png")
    second = await db.create_request(
        user_id=1,
        username="alice",
        mode="i2i",
        prompt_text="banana",
        input_file_path="/tmp/in.png",
        seed=2,
    )
    await db.set_failed(second.id, "error")
    third = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="cherry",
        input_file_path=None,
        seed=3,
    )

    by_user = await db.list_requests(user_id=1)
    assert {item.id for item in by_user} == {first.id, second.id}

    failed = await db.list_requests(status="failed")
    assert [item.id for item in failed] == [second.id]

    queued_t2i = await db.list_requests(mode="t2i", status="queued")
    assert [item.id for item in queued_t2i] == [third.id]

    assert await db.count_requests(user_id=1) == 2
    assert await db.count_requests(status="failed") == 1

    stats = await db.stats()
    assert stats["total"] == 3
    assert stats["users"] == 2
    assert stats["by_status"]["done"] == 1
    assert stats["by_status"]["failed"] == 1
    assert stats["by_status"]["queued"] == 1

    users = await db.list_users()
    assert {user.user_id for user in users} == {1, 2}
    alice = next(user for user in users if user.user_id == 1)
    assert alice.username == "alice"
    assert alice.display_user == "@alice"


async def test_premium_queue_priority_counts(db: Database) -> None:
    normal = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="normal",
        input_file_path=None,
        seed=1,
    )
    premium = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="premium",
        input_file_path=None,
        seed=2,
        is_premium=True,
    )
    counts = await db.queue_counts()
    assert counts == {"premium": 1, "normal": 1}
    assert await db.count_ahead(premium.id) == 0
    assert await db.count_ahead(normal.id) == 1

    await db.set_done(normal.id, "/tmp/normal.png")
    counts = await db.queue_counts()
    assert counts == {"premium": 1, "normal": 0}


async def test_pending_payment_lifecycle(db: Database) -> None:
    pending = await db.create_pending_payment(
        payload="payload-1",
        user_id=1,
        mode="t2i",
        prompt="apple",
        username="alice",
        input_file_path=None,
        amount=0,
    )
    assert pending.status == "pending"
    assert pending.amount == 0

    assert await db.set_pending_choice("payload-1", True, 15) is True
    loaded = await db.get_pending_payment("payload-1")
    assert loaded is not None
    assert loaded.is_premium is True
    assert loaded.amount == 15

    assert await db.mark_pending_paid("payload-1", "charge-1") is True
    loaded = await db.get_pending_payment("payload-1")
    assert loaded is not None
    assert loaded.status == "paid"
    assert loaded.telegram_payment_charge_id == "charge-1"

    await db.attach_pending_request("payload-1", 42)
    loaded = await db.get_pending_payment("payload-1")
    assert loaded is not None
    assert loaded.request_id == 42

    await db.mark_pending_refunded("payload-1")
    loaded = await db.get_pending_payment("payload-1")
    assert loaded is not None
    assert loaded.status == "refunded"

    request = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="paid apple",
        input_file_path=None,
        seed=1,
        payment_charge_id="charge-2",
    )
    assert request.payment_charge_id == "charge-2"
    assert request.payment_refunded is False
    await db.set_refunded(request.id)
    loaded_request = await db.get_request(request.id)
    assert loaded_request is not None
    assert loaded_request.payment_refunded is True


async def test_user_credits_referral_and_daily_accrual(db: Database) -> None:
    user = await db.ensure_user(1, "alice")
    assert user.credits == 3
    assert user.last_daily_credit

    assert await db.add_credits(1, 2) == 5
    assert await db.consume_credit(1) == 4
    assert await db.consume_credit(1) == 3

    friend = await db.ensure_user(2, "bob")
    assert friend.credits == 3
    assert await db.set_referrer(2, 1) is True
    assert await db.set_referrer(2, 1) is False
    assert await db.mark_first_payment(2) is True
    assert await db.mark_first_payment(2) is False

    assert await db.set_age_confirmed(1, True) is True
    user = await db.get_user(1)
    assert user is not None
    assert user.age_confirmed is True
    assert user.age_declined is False
    assert await db.set_age_declined(1) is True
    user = await db.get_user(1)
    assert user is not None
    assert user.age_confirmed is False
    assert user.age_declined is True

    friend = await db.get_user(2)
    assert friend is not None
    assert friend.referrer_id == 1
    assert friend.first_paid_at is not None

    future = msk_today() + timedelta(days=2)
    assert await db.accrue_daily_credits(future) >= 2
    user = await db.get_user(1)
    assert user is not None
    assert user.credits == 5
    assert await db.accrue_daily_credits(future) == 0


async def test_hidden_requests_are_excluded_from_public_queries(
    db: Database,
) -> None:
    visible = await db.create_request(
        user_id=1,
        username="alice",
        mode="t2i",
        prompt_text="visible",
        input_file_path=None,
        seed=1,
    )
    await db.set_done(visible.id, "/tmp/visible.png")
    hidden = await db.create_request(
        user_id=2,
        username="bob",
        mode="t2i",
        prompt_text="hidden",
        input_file_path=None,
        seed=2,
    )
    await db.set_done(hidden.id, "/tmp/hidden.png")
    assert await db.set_hidden(hidden.id, True) is True

    public = await db.list_requests(include_hidden=False)
    assert [item.id for item in public] == [visible.id]
    assert await db.count_requests(include_hidden=False) == 1

    admin = await db.list_requests(include_hidden=True)
    assert {item.id for item in admin} == {visible.id, hidden.id}

    stats = await db.stats(include_hidden=False)
    assert stats["total"] == 1
    assert stats["users"] == 1

    public_users = await db.list_users(include_hidden=False)
    assert {item.user_id for item in public_users} == {1}

    loaded = await db.get_request(hidden.id)
    assert loaded is not None
    assert loaded.hidden is True


async def test_generation_quality_setting(db: Database) -> None:
    user = await db.ensure_user(1, "alice")
    assert user.generation_quality == "standard"
    assert user.generation_steps == 12

    assert await db.set_generation_quality(1, "improved") is True
    assert await db.set_generation_steps(1, 8) is True
    user = await db.get_user(1)
    assert user is not None
    assert user.generation_quality == "improved"
    assert user.generation_steps == 8

    with pytest.raises(ValueError):
        await db.set_generation_quality(1, "unknown")
    with pytest.raises(ValueError):
        await db.set_generation_steps(1, 9)


async def test_user_context_lifecycle(db: Database) -> None:
    assert await db.get_user_context(1) is None

    await db.save_user_context(
        user_id=1,
        base_prompt="A red apple",
        current_prompt="A red apple on a table",
        last_request_id=10,
        last_mode="t2i",
    )
    context = await db.get_user_context(1)
    assert context is not None
    assert context.base_prompt == "A red apple"
    assert context.current_prompt == "A red apple on a table"
    assert context.last_request_id == 10
    assert context.last_mode == "t2i"

    await db.save_user_context(
        user_id=1,
        base_prompt="A red apple",
        current_prompt="A red apple on a blue table",
        last_request_id=11,
        last_mode="t2i",
    )
    context = await db.get_user_context(1)
    assert context is not None
    assert context.base_prompt == "A red apple"
    assert context.current_prompt == "A red apple on a blue table"
    assert context.last_request_id == 11

    assert await db.clear_user_context(1) is True
    assert await db.get_user_context(1) is None
    assert await db.clear_user_context(1) is False


async def test_ensure_user_updates_username_without_resetting_credits(
    db: Database,
) -> None:
    user = await db.ensure_user(1, "alice")
    assert user.credits == 3

    assert await db.add_credits(1, 2) == 5
    user = await db.ensure_user(1, "alice_new")
    assert user.username == "alice_new"
    assert user.credits == 5

    user = await db.ensure_user(1)
    assert user.username == "alice_new"
    assert user.credits == 5