from __future__ import annotations

from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock

import pytest

from codeflash_api.db._queries import Queries
from codeflash_api.db.models import (
    CFAPIKeysRow,
    LLMCallsRow,
    OptimizationErrorsRow,
    SubscriptionsRow,
)


@pytest.fixture(name="pool")
def _pool() -> AsyncMock:
    return AsyncMock()


@pytest.fixture(name="queries")
def _queries(pool: AsyncMock) -> Queries:
    return Queries(pool=pool)


class TestGetApiKeyByHash:
    """Tests for Queries.get_api_key_by_hash."""

    async def test_found(self, queries: Queries, pool: AsyncMock) -> None:
        """
        Matching hash returns an APIKey instance.
        """
        now = datetime.now(tz=timezone.utc)
        pool.fetchrow.return_value = {
            "id": 1,
            "key": "hashed",
            "suffix": "ab12",
            "name": "My Key",
            "user_id": "u1",
            "tier": None,
            "organization_id": None,
            "created_at": now,
            "last_used": None,
        }

        result = await queries.get_api_key_by_hash("hashed")

        assert result is not None
        assert "u1" == result.user_id
        assert 1 == result.id

    async def test_not_found(self, queries: Queries, pool: AsyncMock) -> None:
        """
        Missing hash returns None.
        """
        pool.fetchrow.return_value = None

        result = await queries.get_api_key_by_hash("missing")

        assert result is None


class TestGetSubscription:
    """Tests for Queries.get_subscription."""

    async def test_found(self, queries: Queries, pool: AsyncMock) -> None:
        """
        Existing user_id returns a Subscription.
        """
        pool.fetchrow.return_value = {
            "id": "sub-1",
            "user_id": "u1",
            "plan_type": "FREE",
            "optimizations_used": 10,
            "total_lifetime_optimizations": 100,
            "optimizations_limit": 4000,
            "subscription_status": "active",
            "stripe_customer_id": None,
            "stripe_subscription_id": None,
            "current_period_start": None,
            "current_period_end": None,
            "cancel_at_period_end": False,
            "cancellation_request_date": None,
            "created_at": None,
            "updated_at": None,
        }

        result = await queries.get_subscription("u1")

        assert result is not None
        assert "FREE" == result.plan_type
        assert 10 == result.optimizations_used


class TestIncrementUsage:
    """Tests for Queries.increment_usage."""

    async def test_executes_atomic_update(
        self, queries: Queries, pool: AsyncMock
    ) -> None:
        """
        Increment sends an atomic SQL UPDATE.
        """
        await queries.increment_usage(user_id="u1", cost=10)

        pool.execute.assert_awaited_once()
        sql = pool.execute.call_args[0][0]
        assert "optimizations_used" in sql
        assert "total_lifetime_optimizations" in sql


class TestCreateSubscription:
    """Tests for Queries.create_subscription."""

    async def test_creates_with_defaults(
        self, queries: Queries, pool: AsyncMock
    ) -> None:
        """
        New subscription gets active status and zero usage.
        """
        result = await queries.create_subscription(
            user_id="u1",
            plan_type="FREE",
            optimizations_limit=4000,
        )

        assert "u1" == result.user_id
        assert "FREE" == result.plan_type
        assert 4000 == result.optimizations_limit
        assert 0 == result.optimizations_used
        pool.execute.assert_awaited_once()


class TestDbModels:
    """Tests for database row schema classes."""

    def test_api_key_row_frozen(self) -> None:
        """
        CFAPIKeysRow is immutable.
        """
        row = CFAPIKeysRow(
            id=1,
            key="k",
            suffix="ab",
            name="n",
            created_at=datetime.now(tz=timezone.utc),
        )
        with pytest.raises(AttributeError):
            row.id = 2  # type: ignore[misc]

    def test_llm_calls_row_defaults(self) -> None:
        """
        LLMCallsRow has sensible defaults.
        """
        row = LLMCallsRow(
            trace_id="t1",
            call_type="optimize",
            model_name="gpt-5-mini",
            input_tokens=100,
            output_tokens=50,
        )

        assert 0 == row.cached_tokens
        assert 0.0 == row.cost
        assert 0 == row.latency_ms

    def test_subscriptions_row_defaults(self) -> None:
        """
        SubscriptionsRow defaults match FREE tier.
        """
        row = SubscriptionsRow(
            id="s1",
            user_id="u1",
            plan_type="FREE",
        )

        assert 4000 == row.optimizations_limit
        assert "active" == row.subscription_status
