"""Automated test suite for AI Agent Budget & Token Tracking."""

import unittest

from fastapi.testclient import TestClient

from app.modules.auth.presentation.auth import get_current_user
from app.core.postgres.erp_client import ErpDatabaseClient
from app.main import create_app


class TestAIAgentBudgetService(unittest.TestCase):
    def setUp(self):
        self.app = create_app()
        self.client = TestClient(self.app)
        self.erp_client = ErpDatabaseClient()

    def tearDown(self):
        self.app.dependency_overrides = {}
        with self.erp_client.get_connection() as conn, conn.cursor() as cur:
            cur.execute("DELETE FROM erp_ai_agent_usage_logs WHERE query_preview = 'Test automated budget logging'")
            conn.commit()

    def test_unauthorized_budget_summary(self):
        """Khi chưa đăng nhập, endpoint trả về 401."""
        response = self.client.get("/v1/agents/budget-summary")
        self.assertEqual(response.status_code, 401)

    def test_authorized_budget_summary(self):
        """Khi đã đăng nhập, trả về bảng tổng hợp ngân sách 9 AI agents chuẩn contract."""
        self.app.dependency_overrides[get_current_user] = lambda: {
            "id": "admin-test-id",
            "email": "admin@dscons.vn",
            "role": "superadmin",
        }

        response = self.client.get("/v1/agents/budget-summary")
        self.assertEqual(response.status_code, 200)
        data = response.json()

        self.assertIn("agents", data)
        self.assertIn("total_spent_today_vnd", data)
        self.assertIn("total_spent_week_vnd", data)
        self.assertIn("total_spent_month_vnd", data)
        self.assertIn("total_spent_year_vnd", data)

        agents = data["agents"]
        self.assertGreaterEqual(len(agents), 9)

        # Kiểm tra chi tiết 1 agent (Minh)
        minh = next((a for a in agents if a["agent_code"] == "minh"), None)
        self.assertIsNotNone(minh)
        self.assertEqual(minh["display_name"], "Minh")

        # Kiểm tra 4 chu kỳ Ngày, Tuần, Tháng, Năm
        for period in ["day", "week", "month", "year"]:
            self.assertIn(period, minh)
            p_data = minh[period]
            self.assertIn("spent_vnd", p_data)
            self.assertIn("budget_vnd", p_data)
            self.assertIn("percentage", p_data)
            self.assertIn("tokens", p_data)
            self.assertIn("requests_count", p_data)
            self.assertGreater(p_data["budget_vnd"], 0)

    def test_log_and_recalculate_usage(self):
        """Ghi nhận usage mới và kiểm tra số tiền / token được cộng dồn chính xác."""
        initial_summary = self.erp_client.get_ai_agent_budget_summary(agent_code="minh")
        minh_initial = initial_summary["agents"][0]
        initial_day_spent = minh_initial["day"]["spent_vnd"]
        initial_day_tokens = minh_initial["day"]["tokens"]

        # Log 1 request 500 prompt tokens + 300 completion tokens
        self.erp_client.log_ai_agent_usage(
            agent_code="minh",
            task_type="test_audit",
            prompt_tokens=500,
            completion_tokens=300,
            cost_vnd=40.0,
            user_email="tester@dscons.vn",
            query_preview="Test automated budget logging",
        )

        updated_summary = self.erp_client.get_ai_agent_budget_summary(agent_code="minh")
        minh_updated = updated_summary["agents"][0]

        self.assertAlmostEqual(
            minh_updated["day"]["spent_vnd"], initial_day_spent + 40.0, places=2
        )
        self.assertEqual(minh_updated["day"]["tokens"], initial_day_tokens + 800)


if __name__ == "__main__":
    unittest.main()