import unittest

from collect_daily_usage import coalesce_daily_workspace_rows


class DailyWorkspaceRowsTest(unittest.TestCase):
    def test_multiple_channels_count_workspace_facts_once(self):
        common = {
            "workspace_id": "ws-test",
            "user_email": "a@example.com",
            "message_count": 14,
            "user_message_count": 7,
            "message_categories": {"Ad hoc query": 7},
            "token_usage": {"cost_amount": 5.5, "input_tokens": 1000},
        }
        rows = [
            {**common, "channel": "web", "cost": 5.0},
            {**common, "channel": "image-send", "cost": 0.5},
        ]
        result = coalesce_daily_workspace_rows(rows)
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["cost"], 5.5)
        self.assertEqual(result[0]["message_count"], 14)
        self.assertEqual(result[0]["user_message_count"], 7)
        self.assertEqual(result[0]["token_usage"], common["token_usage"])
        self.assertEqual(result[0]["message_categories"], common["message_categories"])
        self.assertEqual(result[0]["cost_by_channel"], {"web": 5.0, "image-send": 0.5})
        self.assertEqual(result[0]["cost_by_user"], {"a@example.com": 5.5})
        self.assertEqual(rows[0]["cost"], 5.0)

    def test_shared_workspace_preserves_user_costs(self):
        result = coalesce_daily_workspace_rows([
            {"workspace_id": "shared", "user_email": "a", "channel": "web", "cost": 1},
            {"workspace_id": "shared", "user_email": "b", "channel": "web", "cost": 3},
            {"workspace_id": "other", "user_email": "a", "channel": "lark", "cost": 0},
        ])
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0]["cost"], 4)
        self.assertEqual(result[0]["user_email"], "b")
        self.assertEqual(result[0]["cost_by_user"], {"a": 1, "b": 3})
        self.assertEqual(result[1]["workspace_id"], "other")


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