import unittest
from unittest.mock import patch, MagicMock
from pygeai.cli.commands.omni import parse_file
from pygeai.cli.commands import Option
from pygeai.core.common.exceptions import MissingRequirementException, WrongArgumentError


class TestOmniCommands(unittest.TestCase):

    @patch('pygeai.cli.commands.omni.OmniParserClient')
    @patch('pygeai.cli.commands.omni.Console')
    def test_parse_file_basic(self, mock_console, mock_client_class):
        mock_client = MagicMock()
        mock_client_class.return_value = mock_client

        mock_response = {
            "status": "success",
            "parts": [{"text": "Sample text", "type": "Text"}],
            "requestId": "123",
            "usage": {"total_cost": 0.01}
        }
        mock_client.process_file.return_value = mock_response

        option_list = [
            (Option("input", ["--input"], "", True), "test.pdf")
        ]

        parse_file(option_list)

        mock_client.process_file.assert_called_once()
        mock_console.write_stdout.assert_called_once()

    @patch('pygeai.cli.commands.omni.OmniParserClient')
    @patch('pygeai.cli.commands.omni.Console')
    def test_parse_file_with_options(self, mock_console, mock_client_class):
        mock_client = MagicMock()
        mock_client_class.return_value = mock_client

        mock_response = {
            "status": "success",
            "parts": [],
            "requestId": "123"
        }
        mock_client.process_file.return_value = mock_response

        option_list = [
            (Option("input", ["--input"], "", True), "test.pdf"),
            (Option("start_page", ["--start-page"], "", True), "1"),
            (Option("end_page", ["--end-page"], "", True), "5"),
            (Option("hi_res_mode", ["--hi-res"], "", True), "1"),
            (Option("cache", ["--cache"], "", True), "1")
        ]

        parse_file(option_list)

        call_args = mock_client.process_file.call_args
        self.assertEqual(call_args[1]["file_path"], "test.pdf")
        self.assertEqual(call_args[1]["start_page"], 1)
        self.assertEqual(call_args[1]["end_page"], 5)
        self.assertTrue(call_args[1]["hi_res_mode"])
        self.assertTrue(call_args[1]["cache"])

    def test_parse_file_missing_input(self):
        option_list = []

        with self.assertRaises(MissingRequirementException):
            parse_file(option_list)

    def test_parse_file_invalid_page_number(self):
        option_list = [
            (Option("input", ["--input"], "", True), "test.pdf"),
            (Option("start_page", ["--start-page"], "", True), "invalid")
        ]

        with self.assertRaises(WrongArgumentError):
            parse_file(option_list)
