import json
import unittest
from unittest.mock import patch, MagicMock
from consts.consts import Qwen
from errors.errors import BadRequestException, UnauthorizedException, AlgorithmException
from svc.qwen import (
    invoke_qwen_with_text, _get_qwen_api_key, _call_qwen_api, validate_qwen_api_key,
    _get_aliyun_subscription_record, _build_dashscope_euid_header
)


class QwenTests(unittest.TestCase):
    """Qwen model integration unit tests"""

    def setUp(self):
        """Test initialization"""
        self.test_prompt = "Generate SQL for: show all users"
        self.test_api_key = "test-qwen-api-key-12345"
        self.extra_info = {
            "session_id": "test-session-123",
            "query_id": "test-query-456",
            "module": "test-module"
        }

    def tearDown(self):
        """Test cleanup"""
        pass

    # ==================== Test invoke_qwen_with_text ====================

    @patch("svc.qwen._call_qwen_api")
    @patch("svc.qwen._get_qwen_api_key")
    def test_invoke_qwen_with_text_success_json(self, mock_get_api_key, mock_call_api):
        """Test successful Qwen API call returning JSON"""
        # Setup mock
        mock_get_api_key.return_value = self.test_api_key
        mock_call_api.return_value = ('{"sql": "SELECT * FROM users"}', 100, 50)

        # Call function
        result = invoke_qwen_with_text(
            self.test_prompt,
            output_format="json",
            ai_model=Qwen.TURBO,
            extra_info=self.extra_info
        )

        # Verify result
        self.assertIsInstance(result, dict)
        self.assertEqual(result["sql"], "SELECT * FROM users")

        # Verify mock calls
        mock_get_api_key.assert_called_once()
        mock_call_api.assert_called_once()

    @patch("svc.qwen._call_qwen_api")
    @patch("svc.qwen._get_qwen_api_key")
    def test_invoke_qwen_with_text_success_text(self, mock_get_api_key, mock_call_api):
        """Test successful Qwen API call returning text"""
        # Setup mock
        mock_get_api_key.return_value = self.test_api_key
        mock_call_api.return_value = ("This is a response", 100, 50)

        # Call function
        result = invoke_qwen_with_text(
            self.test_prompt,
            output_format="text",
            ai_model=Qwen.PLUS,
            extra_info=self.extra_info
        )

        # Verify result
        self.assertEqual(result, "This is a response")

    @patch("svc.qwen._get_qwen_api_key")
    def test_invoke_qwen_with_text_dry_run(self, mock_get_api_key):
        """Test dry_run mode"""
        # Setup mock
        mock_get_api_key.return_value = self.test_api_key

        # Call function
        result = invoke_qwen_with_text(
            self.test_prompt,
            output_format="json",
            dry_run=True,
            ai_model=Qwen.MAX,
            extra_info=self.extra_info
        )

        # Verify result (dry_run should return True)
        self.assertTrue(result)

    @patch("svc.qwen._call_qwen_api")
    @patch("svc.qwen._get_qwen_api_key")
    @patch("svc.qwen.invoke_qwen_with_text")
    def test_invoke_qwen_with_text_json_repair(self, mock_invoke_self, mock_get_api_key, mock_call_api):
        """Test JSON repair logic"""
        # Setup mock - return corrupted JSON
        mock_get_api_key.return_value = self.test_api_key
        mock_call_api.return_value = ('{"sql": "SELECT * FROM users"', 100, 50)  # Missing closing bracket

        # Since the actual function will recursively call itself, we need to mock the recursive call
        mock_invoke_self.return_value = {"sql": "SELECT * FROM users"}

        # This test verifies that when JSON parsing fails, it will attempt repair
        # The actual implementation will call repair_json and possible recursive calls

    @patch("svc.qwen._call_qwen_api")
    @patch("svc.qwen._get_qwen_api_key")
    def test_invoke_qwen_with_text_invalid_model(self, mock_get_api_key, mock_call_api):
        """Test invalid model ID"""
        # Setup mock
        mock_get_api_key.return_value = self.test_api_key

        # Call function - using invalid model
        with self.assertRaises(BadRequestException):
            invoke_qwen_with_text(
                self.test_prompt,
                output_format="json",
                ai_model="invalid-model",
                extra_info=self.extra_info
            )

    # ==================== Test _get_qwen_api_key ====================

    @patch("svc.qwen.config")
    @patch("svc.qwen.get_local_metrics")
    @patch("svc.qwen.UserPreference.load")
    def test_get_qwen_api_key_user_preference(self, mock_user_pref_load, mock_get_metrics, mock_config):
        """Test user-defined API Key priority"""
        # Setup mock - user has custom key
        mock_config.QWEN_API_KEY = "system-key"
        mock_get_metrics.return_value = {"app_id": 123}
        mock_preference = MagicMock()
        mock_preference.llm_api_key = "user-custom-key"
        mock_user_pref_load.return_value = mock_preference

        # Call function
        result = _get_qwen_api_key(self.extra_info)

        # Verify returns user-defined key
        self.assertEqual(result, "user-custom-key")

    @patch("svc.qwen.config")
    @patch("svc.qwen.get_local_metrics")
    @patch("svc.qwen.UserPreference.load")
    def test_get_qwen_api_key_system_default(self, mock_user_pref_load, mock_get_metrics, mock_config):
        """Test system-level API Key"""
        # Setup mock - user doesn't have custom key
        mock_config.QWEN_API_KEY = "system-key"
        mock_get_metrics.return_value = {"app_id": 123}
        mock_preference = MagicMock()
        mock_preference.llm_api_key = None
        mock_user_pref_load.return_value = mock_preference

        # Call function
        result = _get_qwen_api_key(self.extra_info)

        # Verify returns system-level key
        self.assertEqual(result, "system-key")

    @patch("svc.qwen.config")
    @patch("svc.qwen.get_local_metrics")
    def test_get_qwen_api_key_not_configured(self, mock_get_metrics, mock_config):
        """Test API Key not configured"""
        # Setup mock - no key configured
        mock_config.QWEN_API_KEY = ""
        mock_get_metrics.return_value = {"app_id": 0}

        # Call function - should raise exception
        with self.assertRaises(BadRequestException) as context:
            _get_qwen_api_key(self.extra_info)

        # Verify error message
        self.assertIn("not configured", str(context.exception))

    # ==================== Test _call_qwen_api ====================

    @patch("svc.qwen._build_dashscope_euid_header")
    @patch("svc.qwen.dashscope")
    @patch("svc.qwen.Generation")
    def test_call_qwen_api_success(self, mock_generation_class, mock_dashscope, mock_build_euid):
        """Test successful Qwen API call"""
        # Setup mock response
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.output.choices = [MagicMock()]
        mock_response.output.choices[0].message.content = "API response text"
        mock_response.usage.input_tokens = 100
        mock_response.usage.output_tokens = 50
        mock_generation_class.call.return_value = mock_response
        mock_build_euid.return_value = {}

        # Call function
        result_text, input_tokens, output_tokens = _call_qwen_api(
            "qwen-turbo",
            self.test_prompt,
            [],
            self.test_api_key,
            None
        )

        # Verify result
        self.assertEqual(result_text, "API response text")
        self.assertEqual(input_tokens, 100)
        self.assertEqual(output_tokens, 50)

        # Verify API key passed to Generation.call
        mock_generation_class.call.assert_called_once()
        call_kwargs = mock_generation_class.call.call_args[1]
        self.assertEqual(call_kwargs["api_key"], self.test_api_key)

    @patch("svc.qwen._build_dashscope_euid_header")
    @patch("svc.qwen.dashscope")
    @patch("svc.qwen.Generation")
    def test_call_qwen_api_authentication_error(self, mock_generation_class, mock_dashscope, mock_build_euid):
        """Test authentication failure"""
        # Setup mock response - 401 error
        mock_response = MagicMock()
        mock_response.status_code = 401
        mock_response.message = "Unauthorized"
        mock_generation_class.call.return_value = mock_response
        mock_build_euid.return_value = {}

        # Call function - should raise UnauthorizedException
        with self.assertRaises(UnauthorizedException):
            _call_qwen_api(
                "qwen-turbo",
                self.test_prompt,
                [],
                self.test_api_key,
                None
            )

    @patch("svc.qwen._build_dashscope_euid_header")
    @patch("svc.qwen.dashscope")
    @patch("svc.qwen.Generation")
    def test_call_qwen_api_token_limit_error(self, mock_generation_class, mock_dashscope, mock_build_euid):
        """Test token limit exceeded error"""
        # Setup mock response - Token limit exceeded
        mock_response = MagicMock()
        mock_response.status_code = 400
        mock_response.message = "Token limit exceeded"
        mock_generation_class.call.return_value = mock_response
        mock_build_euid.return_value = {}

        # Call function - should raise BadRequestException
        with self.assertRaises(BadRequestException) as context:
            _call_qwen_api(
                "qwen-turbo",
                self.test_prompt,
                [],
                self.test_api_key,
                None
            )

        # Verify error message
        self.assertIn("extensive tables", str(context.exception))

    @patch("svc.qwen._build_dashscope_euid_header")
    @patch("svc.qwen.dashscope")
    @patch("svc.qwen.Generation")
    def test_call_qwen_api_other_error(self, mock_generation_class, mock_dashscope, mock_build_euid):
        """Test other API errors"""
        # Setup mock response - other errors
        mock_response = MagicMock()
        mock_response.status_code = 500
        mock_response.message = "Internal server error"
        mock_generation_class.call.return_value = mock_response
        mock_build_euid.return_value = {}

        # Call function - should raise AlgorithmException
        with self.assertRaises(AlgorithmException):
            _call_qwen_api(
                "qwen-turbo",
                self.test_prompt,
                [],
                self.test_api_key,
                None
            )

    @patch("svc.qwen.dashscope", None)
    def test_call_qwen_api_missing_dependency(self):
        """Test missing DashScope SDK"""
        # Call function - should raise AlgorithmException (ImportError)
        with self.assertRaises(AlgorithmException) as context:
            _call_qwen_api(
                "qwen-turbo",
                self.test_prompt,
                [],
                self.test_api_key,
                None
            )

        # Verify error message contains installation hint
        self.assertIn("pip install dashscope", str(context.exception))

    # ==================== Test validate_qwen_api_key ====================

    @patch("svc.qwen._call_qwen_api")
    def test_test_qwen_api_key_valid(self, mock_call_api):
        """Test valid API Key"""
        # Setup mock - API call successful
        mock_call_api.return_value = ("test response", 10, 5)

        # Call function
        result = validate_qwen_api_key(self.test_api_key, Qwen.TURBO)

        # Verify result
        self.assertTrue(result)
        mock_call_api.assert_called_once()

    @patch("svc.qwen._call_qwen_api")
    def test_test_qwen_api_key_invalid(self, mock_call_api):
        """Test invalid API Key"""
        # Setup mock - API call failed
        mock_call_api.side_effect = UnauthorizedException("Invalid API key")

        # Call function
        result = validate_qwen_api_key("invalid-key", Qwen.TURBO)

        # Verify result
        self.assertFalse(result)

    @patch("svc.qwen._call_qwen_api")
    def test_test_qwen_api_key_network_error(self, mock_call_api):
        """Test network error"""
        # Setup mock - network error
        mock_call_api.side_effect = Exception("Network error")

        # Call function
        result = validate_qwen_api_key(self.test_api_key, Qwen.TURBO)

        # Verify result
        self.assertFalse(result)


@unittest.skip("Integration tests require kubectl port-forward to billing2 service")
class QwenBillingIntegrationTests(unittest.TestCase):
    """Tests for Qwen billing integration with billing2 service

    To run with local port-forward:
        kubectl port-forward svc/billing2 18888:18888
        pytest tests/test_qwen.py::QwenBillingIntegrationTests -v --no-skip
    """

    def test_get_aliyun_subscription_record_real(self):
        """Test real billing2 API call (requires kubectl port-forward)

        Run: kubectl port-forward svc/billing2 18888:18888
        """
        # Use a known test org_id/tenant_id
        test_org_id = "2130480"

        result = _get_aliyun_subscription_record(test_org_id)

        print(f"\nBilling2 response for org_id={test_org_id}: {result}")

        # Verify response structure
        if result:
            self.assertIn("ali_uid", result)
            self.assertIn("biz_id", result)
            self.assertIn("in_trial", result)
            self.assertIn("is_subscribed", result)

    def test_build_euid_header_subscribed_user(self):
        """Test euid header for subscribed user (requires kubectl port-forward)"""
        # Use a known subscribed org_id
        test_org_id = "2130480"

        headers = _build_dashscope_euid_header(test_org_id)

        print(f"\nEUID headers for org_id={test_org_id}: {headers}")

        # If user is subscribed, headers should contain x-dashscope-euid
        if headers:
            self.assertIn("x-dashscope-euid", headers)
            euid_data = json.loads(headers["x-dashscope-euid"])
            self.assertEqual(euid_data["bizType"], "B2B")
            self.assertEqual(euid_data["moduleType"], "Third-partyproducts")
            self.assertEqual(euid_data["accountType"], "Aliyun")
            self.assertIn("accountId", euid_data)
            self.assertIn("moduleCode", euid_data)
            print(f"EUID data: {euid_data}")

    def test_build_euid_header_not_subscribed(self):
        """Test euid header for non-subscribed user"""
        # Use a non-existent org_id
        test_org_id = "99999999"

        headers = _build_dashscope_euid_header(test_org_id)

        print(f"\nEUID headers for non-subscribed org_id={test_org_id}: {headers}")

        # Non-subscribed user should get empty headers
        self.assertEqual(headers, {})


class QwenBillingMockTests(unittest.TestCase):
    """Mock tests for billing integration"""

    @patch("svc.qwen.requests.get")
    def test_get_aliyun_subscription_record_success(self, mock_get):
        """Test successful billing2 API call"""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "ali_uid": "5293340550807672",
            "biz_id": "5000003055910",
            "in_trial": False,
            "is_subscribed": True
        }
        mock_get.return_value = mock_response

        result = _get_aliyun_subscription_record("12345")

        self.assertIsNotNone(result)
        self.assertEqual(result["biz_id"], "5000003055910")
        self.assertTrue(result["is_subscribed"])

        # Verify correct headers
        mock_get.assert_called_once()
        call_args = mock_get.call_args
        self.assertEqual(call_args[1]["headers"]["X-BILLING-REQUEST-SOURCE"], "ops")

    @patch("svc.qwen.requests.get")
    def test_get_aliyun_subscription_record_not_found(self, mock_get):
        """Test billing2 API 404 response"""
        mock_response = MagicMock()
        mock_response.status_code = 404
        mock_get.return_value = mock_response

        result = _get_aliyun_subscription_record("99999")

        self.assertIsNone(result)

    @patch("svc.qwen.requests.get")
    def test_get_aliyun_subscription_record_timeout(self, mock_get):
        """Test billing2 API timeout"""
        mock_get.side_effect = Exception("Connection timeout")

        result = _get_aliyun_subscription_record("12345")

        self.assertIsNone(result)

    def test_get_aliyun_subscription_record_empty_org_id(self):
        """Test with empty org_id"""
        result = _get_aliyun_subscription_record(None)
        self.assertIsNone(result)

        result = _get_aliyun_subscription_record("")
        self.assertIsNone(result)

    @patch("svc.qwen._get_aliyun_subscription_record")
    def test_build_euid_header_subscribed(self, mock_get_subscription):
        """Test euid header generation for subscribed user"""
        mock_get_subscription.return_value = {
            "ali_uid": "5293340550807672",
            "biz_id": "5000003055910",
            "in_trial": False,
            "is_subscribed": True
        }

        headers = _build_dashscope_euid_header("12345")

        self.assertIn("x-dashscope-euid", headers)
        euid_data = json.loads(headers["x-dashscope-euid"])
        self.assertEqual(euid_data["bizType"], "B2B")
        self.assertEqual(euid_data["moduleType"], "Third-partyproducts")
        self.assertEqual(euid_data["accountType"], "Aliyun")
        self.assertEqual(euid_data["accountId"], "5293340550807672")

    @patch("svc.qwen._get_aliyun_subscription_record")
    def test_build_euid_header_not_subscribed(self, mock_get_subscription):
        """Test euid header generation for non-subscribed user"""
        mock_get_subscription.return_value = {
            "ali_uid": "5293340550807672",
            "biz_id": "5000003055910",
            "in_trial": False,
            "is_subscribed": False
        }

        headers = _build_dashscope_euid_header("12345")

        # Non-subscribed user should get empty headers
        self.assertEqual(headers, {})

    @patch("svc.qwen._get_aliyun_subscription_record")
    def test_build_euid_header_no_subscription_record(self, mock_get_subscription):
        """Test euid header when billing2 returns None"""
        mock_get_subscription.return_value = None

        headers = _build_dashscope_euid_header("12345")

        # No subscription record should get empty headers
        self.assertEqual(headers, {})


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