import json
import unittest
from unittest.mock import patch, mock_open
import urllib3

from aihub.api.run_api import RunApi


class TestRunApi(unittest.TestCase):
  """RunApi unit test stubs"""

  def setUp(self) -> None:
    self.api = RunApi()

  def tearDown(self) -> None:
    pass

  @patch.object(urllib3.PoolManager, 'request')
  def test_get_run_results(self, mock_request) -> None:
    """Test case for get_run_results

        Retrieve run results
        """
    run_id = "run123"
    expected_response = {
        "files": [{
            "original_file_name": "data.pdf",
            "documents": []
        }],
        "has_more": False
    }

    mock_response = urllib3.HTTPResponse(
        body=bytes(json.dumps(expected_response), 'utf-8'), status=200)
    mock_request.return_value = mock_response

    response = self.api.get_run_results(run_id)
    self.assertFalse(response.has_more)
    self.assertEqual(len(response.files), 1)
    self.assertEqual(response.files[0].original_file_name, "data.pdf")

  @patch.object(urllib3.PoolManager, 'request')
  def test_get_run_results_all_flags(self, mock_request) -> None:
    """Test case for get_run_results

        Retrieve run results for all options
        """
    run_id = "run123"
    expected_response = {
        "has_more":
        False,
        "batch_id":
        '1251',
        "files": [{
            "original_file_name":
            "willy.pdf",
            "documents": [{
                "class_name":
                "Wage And Tax Statement",
                "fields": [{
                    "field_name": "field_1",
                    "value": "other",
                    "type": "TEXT",
                    "confidence": {
                        "model": 0.12
                    },
                    "edit_history": [],
                    "source_coordinates": [],
                    "validations": {
                        "valid": True,
                        "alerts": []
                    }
                }, {
                    "field_name":
                    "employers_address_and_ZIP_code",
                    "value":
                    "757 Martin Hill Suite 056 East Kyle SC 74614-1778",
                    "type":
                    "TEXT",
                    "confidence": {
                        "model": 0.70431864
                    },
                    "edit_history": [{
                        "modification_timestamp":
                        "2024-07-01 16:00:00",
                        "user_id":
                        "willy",
                        "modifications": [{
                            "message": "edited to 'X'"
                        }]
                    }],
                    "source_coordinates": [{
                        "top_x": 153.0,
                        "top_y": 573.33,
                        "bottom_x": 191.69999885559082,
                        "bottom_y": 584.08,
                        "page_number": 0
                    }, {
                        "top_x": 205.0,
                        "top_y": 573.33,
                        "bottom_x": 282.39999771118164,
                        "bottom_y": 584.08,
                        "page_number": 0
                    }],
                    "validations": {
                        "valid":
                        False,
                        "alerts": [{
                            "alert_level": "failure",
                            "msg":
                            "Field [Wage And Tax Statement.employers_address_and_ZIP_code] confidence value was lower than threshold 0.95",
                            "blocked": False,
                            "type": "confidence",
                            "locations": []
                        }]
                    }
                }],
                "post_processed_paths": [
                    "instabase-sandbox-org/support-testing.aihubuat-orguser_instabase.com/fs/S3 Test Drive/app-runs/5faf5dc6-dd97-42cb-8b8e-28a63066887e/1716916997868/output/s1_process_files/images/W2_Multi_Sample_Data_input_ADP1_clean_15517-19 (1).pdf_p0.JPEG.jpeg",
                    "instabase-sandbox-org/support-testing.aihubuat-orguser_instabase.com/fs/S3 Test Drive/app-runs/5faf5dc6-dd97-42cb-8b8e-28a63066887e/1716916997868/output/s1_process_files/images/W2_Multi_Sample_Data_input_ADP1_clean_15517-19 (1).pdf_p1.JPEG.jpeg",
                    "instabase-sandbox-org/support-testing.aihubuat-orguser_instabase.com/fs/S3 Test Drive/app-runs/5faf5dc6-dd97-42cb-8b8e-28a63066887e/1716916997868/output/s1_process_files/images/W2_Multi_Sample_Data_input_ADP1_clean_15517-19 (1).pdf_p2.JPEG.jpeg"
                ],
                "review_completed":
                False,
                "class_edit_history": [],
                "validations": {
                    "final_result_pass": False
                },
                "doc_id":
                63333
            }]
        }],
        "review_completed":
        False
    }

    mock_response = urllib3.HTTPResponse(
        body=bytes(json.dumps(expected_response), 'utf-8'), status=200)
    mock_request.return_value = mock_response

    response = self.api.get_run_results(
        run_id,
        include_confidence_scores=True,
        include_source_info=True,
        include_review_results=True,
        include_validation_results=True)
    self.assertFalse(response.has_more)
    self.assertEqual(response.review_completed, False)
    self.assertEqual(len(response.files), 1)
    self.assertEqual(response.files[0].original_file_name, "willy.pdf")
    self.assertEqual(len(response.files[0].documents), 1)
    self.assertEqual(len(response.files[0].documents[0].fields), 2)
    self.assertEqual(
        len(response.files[0].documents[0].post_processed_paths), 3)

    # first field
    first_field = response.files[0].documents[0].fields[0]
    self.assertEqual(first_field.field_name, 'field_1')
    self.assertEqual(first_field.value, 'other')
    self.assertEqual(first_field.confidence.model, 0.12)
    self.assertEqual(first_field.validations.valid, True)
    self.assertEqual(len(first_field.source_coordinates), 0)
    self.assertEqual(len(first_field.edit_history), 0)

    # second field
    second_field = response.files[0].documents[0].fields[1]
    self.assertEqual(second_field.field_name, 'employers_address_and_ZIP_code')
    self.assertEqual(second_field.value,
                     '757 Martin Hill Suite 056 East Kyle SC 74614-1778')
    self.assertEqual(second_field.confidence.model, 0.70431864)
    self.assertEqual(second_field.validations.valid, False)
    self.assertEqual(
        second_field.validations.alerts[0].msg,
        "Field [Wage And Tax Statement.employers_address_and_ZIP_code] confidence value was lower than threshold 0.95"
    )
    self.assertEqual(len(second_field.source_coordinates), 2)
    self.assertEqual(len(second_field.edit_history), 1)
    self.assertEqual(second_field.edit_history[0].modifications[0].message,
                     "edited to 'X'")

  @patch.object(urllib3.PoolManager, 'request')
  def test_get_run_status(self, mock_request) -> None:
    """Test case for get_run_status

        Get the status of a run.
        """
    run_id = "run123"
    expected_response = {"id": run_id, "status": "COMPLETE"}

    mock_response = urllib3.HTTPResponse(
        body=bytes(json.dumps(expected_response), 'utf-8'), status=200)
    mock_request.return_value = mock_response

    response = self.api.get_run_status(run_id)
    self.assertEqual(response.id, run_id)
    self.assertEqual(response.status, "COMPLETE")

  @patch.object(urllib3.PoolManager, 'request')
  def test_run_app(self, mock_request) -> None:
    """Test case for run_app

        Run a given app.
        """
    app_id = "app123"
    batch_id = 123
    expected_response = {"id": "run123", "status": "RUNNING"}

    mock_response = urllib3.HTTPResponse(
        body=bytes(json.dumps(expected_response), 'utf-8'), status=200)
    mock_request.return_value = mock_response

    response = self.api.run_app({'app_id': app_id, 'batch_id': batch_id})
    self.assertEqual(response.id, "run123")
    self.assertEqual(response.status, "RUNNING")

  @patch(
      'builtins.open', new_callable=mock_open, read_data='DUMMY FILE CONTENT')
  @patch.object(urllib3.PoolManager, 'request')
  def test_run_app_sync(self, mock_request, mock_file_open) -> None:
    """Test case for run_app_sync

        Run a given app and return the results.
        """
    app_name = "test_app_sync"
    files = ["data1.pdf"]
    expected_response = {
        "files": [{
            "original_file_name": "data1.pdf",
            "documents": []
        }],
        "has_more": False
    }

    mock_response = urllib3.HTTPResponse(
        body=bytes(json.dumps(expected_response), 'utf-8'), status=200)
    mock_request.return_value = mock_response

    # Assuming the `run_app_sync` method accepts a list of file paths
    # and internally handles file reading or other preprocessing
    response = self.api.run_app_sync(app_name=app_name, files=files)
    self.assertFalse(response.has_more)
    self.assertEqual(len(response.files), 1)
    self.assertEqual(response.files[0].original_file_name, "data1.pdf")


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