"""Tests for calibration components."""

import json
import pytest
import numpy as np
from unittest.mock import AsyncMock, MagicMock, patch

from backend.calibration.vector_learner import (
    learn_model_vector,
    update_model_vector_incremental,
)
from backend.calibration.judge import judge_response, JudgeScore


class TestVectorLearner:
    def test_learn_model_vector_basic(self):
        # Create simple test data
        np.random.seed(42)
        n_samples = 50
        embed_dim = 512
        
        # Random embeddings
        embeddings = np.random.randn(n_samples, embed_dim).astype(np.float32)
        embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
        
        # Random pass/fail labels
        pass_fail = np.random.randint(0, 2, n_samples).astype(np.float32)
        
        vector, accuracy = learn_model_vector(embeddings, pass_fail)
        
        assert vector.shape == (embed_dim,)
        assert 0.0 <= accuracy <= 1.0
        # Vector should be normalized
        assert abs(np.linalg.norm(vector) - 1.0) < 0.01

    def test_learn_model_vector_all_pass(self):
        np.random.seed(42)
        embeddings = np.random.randn(20, 512).astype(np.float32)
        pass_fail = np.ones(20)  # All pass
        
        vector, accuracy = learn_model_vector(embeddings, pass_fail)
        
        # Should return zero vector when all same label
        assert np.allclose(vector, 0.0)

    def test_learn_model_vector_all_fail(self):
        np.random.seed(42)
        embeddings = np.random.randn(20, 512).astype(np.float32)
        pass_fail = np.zeros(20)  # All fail
        
        vector, accuracy = learn_model_vector(embeddings, pass_fail)
        
        assert np.allclose(vector, 0.0)

    def test_learn_model_vector_mismatched_lengths(self):
        embeddings = np.random.randn(20, 512).astype(np.float32)
        pass_fail = np.ones(10)  # Wrong length
        
        with pytest.raises(ValueError, match="Mismatched lengths"):
            learn_model_vector(embeddings, pass_fail)

    def test_learn_model_vector_linear_fallback(self):
        # Test linear method (fallback without sklearn)
        np.random.seed(42)
        embeddings = np.random.randn(30, 512).astype(np.float32)
        pass_fail = np.random.randint(0, 2, 30).astype(np.float32)
        
        vector, accuracy = learn_model_vector(embeddings, pass_fail, method="linear")
        
        assert vector.shape == (512,)
        assert 0.0 <= accuracy <= 1.0

    def test_update_model_vector_incremental(self):
        np.random.seed(42)
        
        # Existing vector
        existing = np.random.randn(512).astype(np.float32)
        existing = existing / np.linalg.norm(existing)
        
        # New data
        new_embeddings = np.random.randn(20, 512).astype(np.float32)
        new_pass_fail = np.random.randint(0, 2, 20).astype(np.float32)
        
        updated = update_model_vector_incremental(
            existing, new_embeddings, new_pass_fail, learning_rate=0.1
        )
        
        assert updated.shape == (512,)
        # Should be normalized
        assert abs(np.linalg.norm(updated) - 1.0) < 0.01
        # Should be different from original
        assert not np.allclose(updated, existing)


class TestJudgeResponse:
    @pytest.mark.asyncio
    async def test_judge_success_openai(self):
        mock_session = MagicMock()
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={
            "choices": [{
                "message": {
                    "content": '{"correctness": 5, "completeness": 4, "following": 5, "reasoning": "Good response"}'
                }
            }]
        })
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock()
        mock_session.post = MagicMock(return_value=mock_response)

        score = await judge_response(
            session=mock_session,
            judge_base_url="https://api.openai.com/v1",
            judge_api_key="sk-test",
            judge_model="gpt-4o",
            judge_provider_type="openai",
            task_prompt="Write hello world",
            model_response="print('Hello, World!')",
            expected_behavior="Should print hello world",
        )

        assert isinstance(score, JudgeScore)
        assert score.correctness == 5.0
        assert score.completeness == 4.0
        assert score.following == 5.0
        assert score.reasoning == "Good response"
        # Composite: 5*0.5 + 4*0.3 + 5*0.2 = 2.5 + 1.2 + 1.0 = 4.7
        assert abs(score.composite - 4.7) < 0.01

    @pytest.mark.asyncio
    async def test_judge_success_anthropic(self):
        mock_session = MagicMock()
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={
            "content": [{
                "type": "text",
                "text": '{"correctness": 4, "completeness": 4, "following": 4, "reasoning": "Solid"}'
            }]
        })
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock()
        mock_session.post = MagicMock(return_value=mock_response)

        score = await judge_response(
            session=mock_session,
            judge_base_url="https://api.anthropic.com",
            judge_api_key="sk-ant-test",
            judge_model="claude-sonnet",
            judge_provider_type="anthropic",
            task_prompt="Test",
            model_response="Response",
            expected_behavior="Expected",
        )

        assert score.correctness == 4.0
        assert score.completeness == 4.0

    @pytest.mark.asyncio
    async def test_judge_api_error(self):
        mock_session = MagicMock()
        mock_response = AsyncMock()
        mock_response.status = 500
        mock_response.text = AsyncMock(return_value="Internal error")
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock()
        mock_session.post = MagicMock(return_value=mock_response)

        score = await judge_response(
            session=mock_session,
            judge_base_url="https://api.openai.com/v1",
            judge_api_key="sk-test",
            judge_model="gpt-4o",
            judge_provider_type="openai",
            task_prompt="Test",
            model_response="Response",
            expected_behavior="Expected",
        )

        # Should return default scores on error
        assert score.correctness == 3.0
        assert score.completeness == 3.0
        assert "500" in score.reasoning

    @pytest.mark.asyncio
    async def test_judge_invalid_json(self):
        mock_session = MagicMock()
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={
            "choices": [{
                "message": {"content": "Not valid JSON at all"}
            }]
        })
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock()
        mock_session.post = MagicMock(return_value=mock_response)

        score = await judge_response(
            session=mock_session,
            judge_base_url="https://api.openai.com/v1",
            judge_api_key="sk-test",
            judge_model="gpt-4o",
            judge_provider_type="openai",
            task_prompt="Test",
            model_response="Response",
            expected_behavior="Expected",
        )

        # Should return default scores
        assert score.correctness == 3.0

    @pytest.mark.asyncio
    async def test_judge_extracts_json_from_text(self):
        mock_session = MagicMock()
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={
            "choices": [{
                "message": {
                    "content": 'Here is my evaluation:\n{"correctness": 5, "completeness": 5, "following": 5, "reasoning": "Perfect"}\nDone.'
                }
            }]
        })
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock()
        mock_session.post = MagicMock(return_value=mock_response)

        score = await judge_response(
            session=mock_session,
            judge_base_url="https://api.openai.com/v1",
            judge_api_key="sk-test",
            judge_model="gpt-4o",
            judge_provider_type="openai",
            task_prompt="Test",
            model_response="Response",
            expected_behavior="Expected",
        )

        assert score.correctness == 5.0
        assert score.reasoning == "Perfect"


class TestJudgeScore:
    def test_judge_score_composite(self):
        score = JudgeScore(
            correctness=5.0,
            completeness=4.0,
            following=3.0,
            composite=4.3,
            reasoning="Test",
        )
        # Verify composite calculation: 5*0.5 + 4*0.3 + 3*0.2 = 2.5 + 1.2 + 0.6 = 4.3
        assert score.composite == 4.3
