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


class TestOmniParserResponses(unittest.TestCase):

    def test_omni_parser_part(self):
        part = OmniParserPart(
            page=1,
            type="NarrativeText",
            text="Sample text",
            mid_time=None
        )

        self.assertEqual(part.page, 1)
        self.assertEqual(part.type, "NarrativeText")
        self.assertEqual(part.text, "Sample text")
        self.assertIsNone(part.mid_time)

    def test_omni_parser_part_video(self):
        part = OmniParserPart(
            text="Spoken words",
            mid_time=45
        )

        self.assertEqual(part.text, "Spoken words")
        self.assertEqual(part.mid_time, 45)
        self.assertIsNone(part.page)

    def test_omni_parser_usage(self):
        usage = OmniParserUsage(
            total_cost=0.05,
            total_tokens=100,
            currency="USD"
        )

        self.assertEqual(usage.total_cost, 0.05)
        self.assertEqual(usage.total_tokens, 100)
        self.assertEqual(usage.currency, "USD")

    def test_omni_parser_response(self):
        parts = [
            OmniParserPart(page=1, text="Text 1", type="Title"),
            OmniParserPart(page=2, text="Text 2", type="NarrativeText")
        ]
        usage = OmniParserUsage(total_cost=0.01, total_tokens=50, currency="USD")

        response = OmniParserResponse(
            status="success",
            parts=parts,
            request_id="test-123",
            usage=usage
        )

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