import unittest
from pygeai.core.omni.mappers import OmniParserResponseMapper
from pygeai.core.omni.responses import OmniParserResponse, OmniParserPart, OmniParserUsage


class TestOmniParserResponseMapper(unittest.TestCase):

    def test_map_to_omni_parser_response(self):
        data = {
            "status": "success",
            "parts": [
                {"page": 1, "type": "NarrativeText", "text": "Sample text"},
                {"page": 2, "type": "Title", "text": "Chapter 1"}
            ],
            "requestId": "test-123",
            "usage": {
                "total_cost": 0.05,
                "total_tokens": 100,
                "currency": "USD"
            }
        }

        result = OmniParserResponseMapper.map_to_omni_parser_response(data)

        self.assertIsInstance(result, OmniParserResponse)
        self.assertEqual(result.status, "success")
        self.assertEqual(result.request_id, "test-123")
        self.assertEqual(len(result.parts), 2)
        self.assertIsInstance(result.usage, OmniParserUsage)

    def test_map_to_part_list(self):
        data = [
            {"page": 1, "type": "NarrativeText", "text": "Sample text"},
            {"midTime": 30, "type": "dialogue", "text": "Spoken words"}
        ]

        result = OmniParserResponseMapper.map_to_part_list(data)

        self.assertEqual(len(result), 2)
        self.assertIsInstance(result[0], OmniParserPart)
        self.assertEqual(result[0].page, 1)
        self.assertEqual(result[1].mid_time, 30)

    def test_map_to_part(self):
        data = {"page": 5, "type": "ListItem", "text": "Item 1"}

        result = OmniParserResponseMapper.map_to_part(data)

        self.assertIsInstance(result, OmniParserPart)
        self.assertEqual(result.page, 5)
        self.assertEqual(result.type, "ListItem")
        self.assertEqual(result.text, "Item 1")

    def test_map_to_usage(self):
        data = {
            "total_cost": 0.05,
            "total_tokens": 100,
            "prompt_tokens": 80,
            "completion_tokens": 20,
            "currency": "USD"
        }

        result = OmniParserResponseMapper.map_to_usage(data)

        self.assertIsInstance(result, OmniParserUsage)
        self.assertEqual(result.total_cost, 0.05)
        self.assertEqual(result.total_tokens, 100)
        self.assertEqual(result.currency, "USD")

    def test_map_to_usage_none(self):
        result = OmniParserResponseMapper.map_to_usage(None)
        self.assertIsNone(result)
