import pytest
from datetime import UTC, datetime
from unittest.mock import MagicMock, AsyncMock

from app.services.ai_engine.base import AIEngineAdapter, AIQueryResult, EngineType
from app.services.usage_tracker import UsageTracker


def _make_result_with_tokens(
    engine_type: EngineType,
    query: str = "test query",
    input_tokens: int = 100,
    output_tokens: int = 200,
    has_brand: bool = False,
) -> AIQueryResult:
    return AIQueryResult(
        engine_type=engine_type,
        query=query,
        raw_response="some response with content",
        citations=[],
        has_brand_citation=has_brand,
        has_competitor_citation=False,
        brand_context="brand context" if has_brand else None,
        competitor_contexts=[],
        response_time_ms=100,
        timestamp=datetime.now(UTC),
        input_tokens=input_tokens,
        output_tokens=output_tokens,
    )


class _MockAdapter(AIEngineAdapter):
    def __init__(self, engine_type: EngineType, result: AIQueryResult | None = None):
        super().__init__(api_key="test-key")
        self._engine_type = engine_type
        self._result = result

    async def query(self, query: str, brand_name: str, competitor_names: list[str] | None = None) -> AIQueryResult:
        return self._result

    def get_engine_type(self) -> EngineType:
        return self._engine_type

    def _get_env_key(self) -> str | None:
        return None


class TestUsageRecordingOnQuery:
    """测试查询后自动记录用量"""

    @pytest.mark.asyncio
    async def test_query_single_records_usage(self):
        from app.services.ai_engine.batch_query import BatchQueryService
        from app.services.usage_recorder import UsageRecorder

        result = _make_result_with_tokens(EngineType.CHATGPT, input_tokens=100, output_tokens=200)
        adapters = {"chatgpt": _MockAdapter(EngineType.CHATGPT, result=result)}
        service = BatchQueryService(adapters)
        service.set_user_context("user123", "brand456")

        await service.query_single(EngineType.CHATGPT, "test query", "BrandX")

        summary = service.get_usage_summary()
        assert summary.total_queries == 1
        assert summary.total_input_tokens == 100
        assert summary.total_output_tokens == 200
        assert summary.total_cost > 0

    @pytest.mark.asyncio
    async def test_batch_query_records_usage_for_all(self):
        from app.services.ai_engine.batch_query import BatchQueryService

        r1 = _make_result_with_tokens(EngineType.CHATGPT, input_tokens=100, output_tokens=200)
        r2 = _make_result_with_tokens(EngineType.DEEPSEEK, input_tokens=150, output_tokens=300)
        adapters = {
            "chatgpt": _MockAdapter(EngineType.CHATGPT, result=r1),
            "deepseek": _MockAdapter(EngineType.DEEPSEEK, result=r2),
        }
        service = BatchQueryService(adapters)
        service.set_user_context("user123")

        await service.query_batch(
            [EngineType.CHATGPT, EngineType.DEEPSEEK],
            "test query",
            "BrandX",
        )

        summary = service.get_usage_summary()
        assert summary.total_queries == 2
        assert summary.total_input_tokens == 250
        assert summary.total_output_tokens == 500

    @pytest.mark.asyncio
    async def test_no_user_context_no_recording(self):
        from app.services.ai_engine.batch_query import BatchQueryService

        result = _make_result_with_tokens(EngineType.CHATGPT, input_tokens=100, output_tokens=200)
        adapters = {"chatgpt": _MockAdapter(EngineType.CHATGPT, result=result)}
        service = BatchQueryService(adapters)

        await service.query_single(EngineType.CHATGPT, "test query", "BrandX")

        summary = service.get_usage_summary()
        assert summary.total_queries == 0


class TestUsageRecorderCostCalculation:
    """测试不同引擎记录正确的成本"""

    def test_chatgpt_cost_calculation(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        cost = recorder.calculate_cost("chatgpt", 1000, 2000)
        expected_input = (1000 / 1_000_000) * 1.0
        expected_output = (2000 / 1_000_000) * 4.0
        expected = round(expected_input + expected_output, 6)
        assert cost == expected

    def test_deepseek_cost_calculation(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        cost = recorder.calculate_cost("deepseek", 1000, 2000)
        expected_input = (1000 / 1_000_000) * 0.25
        expected_output = (2000 / 1_000_000) * 6.0
        expected = round(expected_input + expected_output, 6)
        assert cost == expected

    def test_kimi_cost_calculation(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        cost = recorder.calculate_cost("kimi", 1000, 2000)
        expected_input = (1000 / 1_000_000) * 12.0
        expected_output = (2000 / 1_000_000) * 12.0
        expected = round(expected_input + expected_output, 6)
        assert cost == expected

    def test_unknown_engine_zero_cost(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        cost = recorder.calculate_cost("unknown_engine", 1000, 2000)
        assert cost == 0.0

    def test_record_with_cost(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        recorder.record(
            user_id="user123",
            brand_id="brand456",
            engine_type="chatgpt",
            query="test query",
            input_tokens=1000,
            output_tokens=2000,
        )

        summary = tracker.get_summary(user_id="user123")
        assert summary.total_queries == 1
        assert summary.total_input_tokens == 1000
        assert summary.total_output_tokens == 2000
        assert summary.total_cost > 0


class TestUsageSummaryCalculation:
    """测试用量汇总计算正确"""

    def test_summary_by_engine(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        recorder.record("user1", "brand1", "chatgpt", "q1", 100, 200)
        recorder.record("user1", "brand1", "chatgpt", "q2", 150, 250)
        recorder.record("user1", "brand1", "deepseek", "q3", 200, 300)

        summary = tracker.get_summary(user_id="user1")
        assert summary.by_engine["chatgpt"]["queries"] == 2
        assert summary.by_engine["chatgpt"]["input_tokens"] == 250
        assert summary.by_engine["deepseek"]["queries"] == 1

    def test_summary_by_day(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        recorder.record("user1", "brand1", "chatgpt", "q1", 100, 200)
        recorder.record("user1", "brand1", "chatgpt", "q2", 150, 250)

        summary = tracker.get_summary(user_id="user1")
        today = datetime.now(UTC).strftime("%Y-%m-%d")
        assert today in summary.by_day
        assert summary.by_day[today]["queries"] == 2

    def test_summary_filter_by_brand(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        recorder.record("user1", "brand1", "chatgpt", "q1", 100, 200)
        recorder.record("user1", "brand2", "chatgpt", "q2", 150, 250)

        summary = tracker.get_summary(user_id="user1", brand_id="brand1")
        assert summary.total_queries == 1
        assert summary.total_input_tokens == 100


class TestQuotaCheck:
    """测试配额检查触发"""

    def test_quota_check_warning(self):
        from app.services.usage_recorder import UsageRecorder

        tracker = UsageTracker()
        recorder = UsageRecorder(tracker)

        cost = recorder.calculate_cost("chatgpt", 1000000, 2000000)
        recorder.record("user1", "", "chatgpt", "q1", 1000000, 2000000)
        quota = tracker.check_quota("user1", monthly_limit=cost / 2)

        assert quota["status"] == "exceeded"
        assert quota["used"] > 0

    def test_quota_check_ok(self):
        tracker = UsageTracker()
        quota = tracker.check_quota("new_user", monthly_limit=100.0)
        assert quota["status"] == "ok"
        assert quota["used"] == 0


class TestBatchQueryServiceUsageIntegration:
    """测试BatchQueryService与用量记录的集成"""

    @pytest.mark.asyncio
    async def test_service_has_usage_recorder(self):
        from app.services.ai_engine.batch_query import BatchQueryService

        adapters = {"chatgpt": _MockAdapter(EngineType.CHATGPT)}
        service = BatchQueryService(adapters)
        assert hasattr(service, "_recorder")
        assert hasattr(service, "_tracker")

    @pytest.mark.asyncio
    async def test_set_user_context_stores_context(self):
        from app.services.ai_engine.batch_query import BatchQueryService

        adapters = {"chatgpt": _MockAdapter(EngineType.CHATGPT)}
        service = BatchQueryService(adapters)
        service.set_user_context("user123", "brand456")

        assert service._user_id == "user123"
        assert service._brand_id == "brand456"

    @pytest.mark.asyncio
    async def test_get_usage_summary_returns_summary(self):
        from app.services.ai_engine.batch_query import BatchQueryService

        adapters = {"chatgpt": _MockAdapter(EngineType.CHATGPT)}
        service = BatchQueryService(adapters)

        summary = service.get_usage_summary()
        assert hasattr(summary, "total_queries")
        assert hasattr(summary, "total_input_tokens")
        assert hasattr(summary, "total_output_tokens")
        assert hasattr(summary, "total_cost")


class TestAIQueryResultTokens:
    """测试AIQueryResult包含token统计字段"""

    def test_result_has_token_fields(self):
        result = AIQueryResult(
            engine_type=EngineType.CHATGPT,
            query="test",
            raw_response="response",
            citations=[],
            has_brand_citation=False,
            has_competitor_citation=False,
            brand_context=None,
            competitor_contexts=[],
            response_time_ms=100,
            timestamp=datetime.now(UTC),
            input_tokens=100,
            output_tokens=200,
        )
        assert result.input_tokens == 100
        assert result.output_tokens == 200
        assert result.total_tokens == 300

    def test_result_token_fields_default_to_zero(self):
        result = AIQueryResult(
            engine_type=EngineType.CHATGPT,
            query="test",
            raw_response="response",
            citations=[],
            has_brand_citation=False,
            has_competitor_citation=False,
            brand_context=None,
            competitor_contexts=[],
            response_time_ms=100,
            timestamp=datetime.now(UTC),
        )
        assert result.input_tokens == 0
        assert result.output_tokens == 0
        assert result.total_tokens == 0
