"""Tests for UsageRecord model."""
import uuid
from datetime import datetime, timezone, timedelta

import pytest
from sqlalchemy import select, func, and_
from sqlalchemy.dialects.postgresql import insert as pg_insert

from app.models.usage_record import UsageRecord


class TestUsageRecordModel:
    """Test cases for UsageRecord model."""

    @pytest.mark.asyncio
    async def test_usage_record_create(self, async_session, test_user):
        """Test creating a new usage record."""
        record = UsageRecord(
            user_id=test_user.id,
            engine_type="chatgpt",
            query="What is SEO optimization?",
            input_tokens=100,
            output_tokens=200,
            cost=0.015,
            extra_data={"model": "gpt-4"},
        )
        async_session.add(record)
        await async_session.commit()
        await async_session.refresh(record)

        assert record.id is not None
        assert record.user_id == test_user.id
        assert record.engine_type == "chatgpt"
        assert record.query == "What is SEO optimization?"
        assert record.input_tokens == 100
        assert record.output_tokens == 200
        assert record.cost == 0.015
        assert record.extra_data == {"model": "gpt-4"}
        assert record.timestamp is not None
        assert record.created_at is not None

    @pytest.mark.asyncio
    async def test_usage_record_default_values(self, async_session, test_user):
        """Test usage record default values."""
        record = UsageRecord(
            user_id=test_user.id,
            engine_type="kimi",
            query="Test query",
        )
        async_session.add(record)
        await async_session.commit()
        await async_session.refresh(record)

        assert record.input_tokens == 0
        assert record.output_tokens == 0
        assert record.cost == 0.0
        assert record.extra_data == {}
        assert record.brand_id is None

    @pytest.mark.asyncio
    async def test_usage_record_query_by_user_id(self, async_session, test_user):
        """Test querying usage records by user ID."""
        user_id = test_user.id
        for i in range(3):
            record = UsageRecord(
                user_id=user_id,
                engine_type="deepseek",
                query=f"Test query {i}",
                cost=float(i),
            )
            async_session.add(record)
        await async_session.commit()

        result = await async_session.execute(
            select(UsageRecord).where(UsageRecord.user_id == user_id)
        )
        records = result.scalars().all()

        assert len(records) == 3

    @pytest.mark.asyncio
    async def test_usage_record_query_by_user_and_engine(self, async_session, test_user):
        """Test querying usage records by user ID and engine type."""
        user_id = test_user.id
        record1 = UsageRecord(
            user_id=user_id,
            engine_type="chatgpt",
            query="Test chatgpt",
            cost=1.0,
        )
        record2 = UsageRecord(
            user_id=user_id,
            engine_type="kimi",
            query="Test kimi",
            cost=2.0,
        )
        async_session.add(record1)
        async_session.add(record2)
        await async_session.commit()

        result = await async_session.execute(
            select(UsageRecord).where(
                and_(
                    UsageRecord.user_id == user_id,
                    UsageRecord.engine_type == "chatgpt"
                )
            )
        )
        records = result.scalars().all()

        assert len(records) == 1
        assert records[0].engine_type == "chatgpt"

    @pytest.mark.asyncio
    async def test_usage_record_query_by_time_range(self, async_session, test_user):
        """Test querying usage records by time range."""
        user_id = test_user.id
        now = datetime.now(timezone.utc)
        
        old_record = UsageRecord(
            user_id=user_id,
            engine_type="gemini",
            query="Old query",
            cost=1.0,
            timestamp=now - timedelta(days=10),
        )
        new_record = UsageRecord(
            user_id=user_id,
            engine_type="gemini",
            query="New query",
            cost=2.0,
            timestamp=now,
        )
        async_session.add(old_record)
        async_session.add(new_record)
        await async_session.commit()

        cutoff = now - timedelta(days=7)
        result = await async_session.execute(
            select(UsageRecord).where(
                and_(
                    UsageRecord.user_id == user_id,
                    UsageRecord.timestamp >= cutoff
                )
            )
        )
        records = result.scalars().all()

        assert len(records) == 1
        assert records[0].query == "New query"

    @pytest.mark.asyncio
    async def test_usage_record_aggregate_by_user(self, async_session, test_user):
        """Test aggregating usage records by user."""
        user_id = test_user.id
        records_data = [
            {"engine": "chatgpt", "input_tokens": 100, "output_tokens": 200, "cost": 0.01},
            {"engine": "kimi", "input_tokens": 150, "output_tokens": 300, "cost": 0.02},
            {"engine": "chatgpt", "input_tokens": 200, "output_tokens": 400, "cost": 0.03},
        ]
        for data in records_data:
            record = UsageRecord(
                user_id=user_id,
                engine_type=data["engine"],
                query=f"Query for {data['engine']}",
                input_tokens=data["input_tokens"],
                output_tokens=data["output_tokens"],
                cost=data["cost"],
            )
            async_session.add(record)
        await async_session.commit()

        result = await async_session.execute(
            select(
                UsageRecord.engine_type,
                func.count(UsageRecord.id).label("count"),
                func.sum(UsageRecord.input_tokens).label("total_input"),
                func.sum(UsageRecord.output_tokens).label("total_output"),
                func.sum(UsageRecord.cost).label("total_cost"),
            ).where(UsageRecord.user_id == user_id).group_by(UsageRecord.engine_type)
        )
        aggregates = result.all()

        assert len(aggregates) == 2
        
        chatgpt_agg = next(a for a in aggregates if a.engine_type == "chatgpt")
        assert chatgpt_agg.count == 2
        assert chatgpt_agg.total_input == 300
        assert chatgpt_agg.total_output == 600
        assert chatgpt_agg.total_cost == 0.04

    @pytest.mark.asyncio
    async def test_usage_record_aggregate_by_day(self, async_session, test_user):
        """Test aggregating usage records by day."""
        user_id = test_user.id
        now = datetime.now(timezone.utc)
        today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
        
        record1 = UsageRecord(
            user_id=user_id,
            engine_type="qwen",
            query="Query today 1",
            cost=1.0,
            timestamp=today_start + timedelta(hours=10),
        )
        record2 = UsageRecord(
            user_id=user_id,
            engine_type="qwen",
            query="Query today 2",
            cost=2.0,
            timestamp=today_start + timedelta(hours=14),
        )
        record3 = UsageRecord(
            user_id=user_id,
            engine_type="qwen",
            query="Query yesterday",
            cost=3.0,
            timestamp=today_start - timedelta(days=1),
        )
        async_session.add_all([record1, record2, record3])
        await async_session.commit()

        result = await async_session.execute(
            select(
                func.date(UsageRecord.timestamp).label("date"),
                func.count(UsageRecord.id).label("count"),
                func.sum(UsageRecord.cost).label("total_cost"),
            ).where(
                UsageRecord.user_id == user_id
            ).group_by(func.date(UsageRecord.timestamp))
        )
        aggregates = result.all()

        assert len(aggregates) == 2

    @pytest.mark.asyncio
    async def test_usage_record_with_brand_id(self, async_session, test_user):
        """Test usage record with brand association."""
        brand_id = uuid.uuid4()
        record = UsageRecord(
            user_id=test_user.id,
            brand_id=brand_id,
            engine_type="wenxin",
            query="Brand query",
            cost=5.0,
        )
        async_session.add(record)
        await async_session.commit()
        await async_session.refresh(record)

        assert record.brand_id == brand_id

    @pytest.mark.asyncio
    async def test_usage_record_index_user_engine(self, async_session, test_user):
        """Test composite index on user_id and engine_type."""
        user_id = test_user.id
        for i in range(5):
            record = UsageRecord(
                user_id=user_id,
                engine_type="doubao",
                query=f"Doubao query {i}",
                cost=float(i),
            )
            async_session.add(record)
        await async_session.commit()

        result = await async_session.execute(
            select(UsageRecord).where(
                and_(
                    UsageRecord.user_id == user_id,
                    UsageRecord.engine_type == "doubao"
                )
            )
        )
        records = result.scalars().all()

        assert len(records) == 5

    @pytest.mark.asyncio
    async def test_usage_record_update(self, async_session, test_user):
        """Test updating usage record fields."""
        record = UsageRecord(
            user_id=test_user.id,
            engine_type="xinghuo",
            query="Original query",
            cost=1.0,
        )
        async_session.add(record)
        await async_session.commit()

        record.cost = 10.0
        record.extra_data = {"updated": True}
        await async_session.commit()
        await async_session.refresh(record)

        assert record.cost == 10.0
        assert record.extra_data == {"updated": True}

    @pytest.mark.asyncio
    async def test_usage_record_delete(self, async_session, test_user):
        """Test deleting a usage record."""
        record = UsageRecord(
            user_id=test_user.id,
            engine_type="yuanbao",
            query="Delete me",
            cost=1.0,
        )
        async_session.add(record)
        await async_session.commit()
        record_id = record.id

        await async_session.delete(record)
        await async_session.commit()

        result = await async_session.execute(
            select(UsageRecord).where(UsageRecord.id == record_id)
        )
        deleted_record = result.scalar_one_or_none()

        assert deleted_record is None

    @pytest.mark.asyncio
    async def test_usage_record_timestamps(self, async_session, test_user):
        """Test usage record timestamp fields."""
        record = UsageRecord(
            user_id=test_user.id,
            engine_type="perplexity",
            query="Timestamp test",
            cost=1.0,
        )
        async_session.add(record)
        await async_session.commit()
        await async_session.refresh(record)

        assert record.timestamp is not None
        assert record.created_at is not None
        assert isinstance(record.timestamp, datetime)
        assert isinstance(record.created_at, datetime)

    @pytest.mark.asyncio
    async def test_usage_record_query_multiple_users(self, async_session, test_user):
        """Test querying records for multiple users."""
        other_user_id = uuid.uuid4()
        
        user1_record = UsageRecord(
            user_id=test_user.id,
            engine_type="chatgpt",
            query="User 1 query",
            cost=1.0,
        )
        user2_record = UsageRecord(
            user_id=other_user_id,
            engine_type="kimi",
            query="User 2 query",
            cost=2.0,
        )
        async_session.add(user1_record)
        async_session.add(user2_record)
        await async_session.commit()

        result_user1 = await async_session.execute(
            select(UsageRecord).where(UsageRecord.user_id == test_user.id)
        )
        result_user2 = await async_session.execute(
            select(UsageRecord).where(UsageRecord.user_id == other_user_id)
        )

        assert len(result_user1.scalars().all()) == 1
        assert len(result_user2.scalars().all()) == 1

    @pytest.mark.asyncio
    async def test_usage_record_empty_query_field(self, async_session, test_user):
        """Test usage record with empty query field."""
        record = UsageRecord(
            user_id=test_user.id,
            engine_type="deepseek",
            query="",
            cost=0.0,
        )
        async_session.add(record)
        await async_session.commit()
        await async_session.refresh(record)

        assert record.query == ""

    @pytest.mark.asyncio
    async def test_usage_record_large_metadata(self, async_session, test_user):
        """Test usage record with large metadata dict."""
        large_metadata = {
            "model": "gpt-4-turbo",
            "temperature": 0.7,
            "max_tokens": 1000,
            "system_prompt": "You are a helpful assistant." * 100,
            "nested": {"key": {"deep": {"value": [1, 2, 3]}}},
        }
        record = UsageRecord(
            user_id=test_user.id,
            engine_type="chatgpt",
            query="Large metadata test",
            extra_data=large_metadata,
        )
        async_session.add(record)
        await async_session.commit()
        await async_session.refresh(record)

        assert record.extra_data == large_metadata
