import os
import sys
from unittest.mock import MagicMock

sys.path.insert(
    0, os.path.abspath("../../../../..")
)  # Adds the parent directory to the system path

from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig
from litellm.types.utils import EmbeddingResponse


class TestCohereEmbeddingV1Transform:
    def setup_method(self):
        self.config = CohereEmbeddingConfig()
        self.model = "embed-english-v3.0"
        self.logging_obj = MagicMock()
        self.encoding = MagicMock()
        # Mock the encoding to return a fixed token count
        self.encoding.encode = MagicMock(return_value=[1, 2, 3, 4, 5])

    def test_transform_response_regular_embeddings(self):
        """Test that regular embeddings are correctly transformed"""
        # Mock httpx.Response
        mock_response = MagicMock()
        response_json = {
            "embeddings": [
                [0.1, 0.2, 0.3],
                [0.4, 0.5, 0.6],
            ],
            "meta": {"billed_units": {"input_tokens": 10}},
        }
        mock_response.json = MagicMock(return_value=response_json)

        input_data = ["test text 1", "test text 2"]
        data = {"texts": input_data, "input_type": "search_query"}
        model_response = EmbeddingResponse()

        result = self.config._transform_response(
            response=mock_response,
            api_key="test-api-key",
            logging_obj=self.logging_obj,
            data=data,
            model_response=model_response,
            model=self.model,
            encoding=self.encoding,
            input=input_data,
        )

        # Verify the response structure
        assert result.object == "list"
        assert result.model == self.model
        assert len(result.data) == 2

        # Verify each embedding object
        assert result.data[0]["object"] == "embedding"
        assert result.data[0]["index"] == 0
        assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
        assert "type" not in result.data[0]

        assert result.data[1]["object"] == "embedding"
        assert result.data[1]["index"] == 1
        assert result.data[1]["embedding"] == [0.4, 0.5, 0.6]
        assert "type" not in result.data[1]

        # Verify usage
        assert result.usage is not None
        assert result.usage.prompt_tokens == 10
        assert result.usage.total_tokens == 10
        assert result.usage.completion_tokens == 0

    def test_transform_response_embeddings_by_type(self):
        """Test that embeddings_by_type are correctly transformed"""
        # Mock httpx.Response
        mock_response = MagicMock()
        response_json = {
            "response_type": "embeddings_by_type",
            "embeddings": {
                "float": [
                    [0.1, 0.2, 0.3],
                    [0.4, 0.5, 0.6],
                ],
                "int8": [
                    [1, 2, 3],
                    [4, 5, 6],
                ],
            },
            "meta": {"billed_units": {"input_tokens": 10}},
        }
        mock_response.json = MagicMock(return_value=response_json)

        input_data = ["test text 1", "test text 2"]
        data = {
            "texts": input_data,
            "input_type": "search_query",
            "embedding_types": ["float", "int8"],
        }
        model_response = EmbeddingResponse()

        result = self.config._transform_response(
            response=mock_response,
            api_key="test-api-key",
            logging_obj=self.logging_obj,
            data=data,
            model_response=model_response,
            model=self.model,
            encoding=self.encoding,
            input=input_data,
        )

        # Verify the response structure
        assert result.object == "list"
        assert result.model == self.model
        assert len(result.data) == 4  # 2 texts * 2 embedding types

        # Verify float embeddings
        assert result.data[0]["object"] == "embedding"
        assert result.data[0]["index"] == 0
        assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
        assert result.data[0]["type"] == "float"

        assert result.data[1]["object"] == "embedding"
        assert result.data[1]["index"] == 1
        assert result.data[1]["embedding"] == [0.4, 0.5, 0.6]
        assert result.data[1]["type"] == "float"

        # Verify int8 embeddings
        assert result.data[2]["object"] == "embedding"
        assert result.data[2]["index"] == 0
        assert result.data[2]["embedding"] == [1, 2, 3]
        assert result.data[2]["type"] == "int8"

        assert result.data[3]["object"] == "embedding"
        assert result.data[3]["index"] == 1
        assert result.data[3]["embedding"] == [4, 5, 6]
        assert result.data[3]["type"] == "int8"

        # Verify usage
        assert result.usage is not None
        assert result.usage.prompt_tokens == 10
        assert result.usage.total_tokens == 10
        assert result.usage.completion_tokens == 0

    def test_transform_response_with_image_tokens(self):
        """Test that image token billing is correctly handled"""
        # Mock httpx.Response
        mock_response = MagicMock()
        response_json = {
            "embeddings": [
                [0.1, 0.2, 0.3],
            ],
            "meta": {"billed_units": {"input_tokens": 5, "images": 100}},
        }
        mock_response.json = MagicMock(return_value=response_json)

        input_data = ["test image"]
        data = {"images": input_data, "input_type": "image"}
        model_response = EmbeddingResponse()

        result = self.config._transform_response(
            response=mock_response,
            api_key="test-api-key",
            logging_obj=self.logging_obj,
            data=data,
            model_response=model_response,
            model=self.model,
            encoding=self.encoding,
            input=input_data,
        )

        # Verify usage includes both text and image tokens
        assert result.usage is not None
        assert result.usage.prompt_tokens == 105  # 5 text + 100 image
        assert result.usage.total_tokens == 105
        assert result.usage.completion_tokens == 0
        assert result.usage.prompt_tokens_details is not None
        assert result.usage.prompt_tokens_details.text_tokens == 5
        assert result.usage.prompt_tokens_details.image_tokens == 100

    def test_transform_response_fallback_token_counting(self):
        """Test that token counting falls back to encoding when billed_units not present"""
        # Mock httpx.Response
        mock_response = MagicMock()
        response_json = {
            "embeddings": [
                [0.1, 0.2, 0.3],
            ],
            "meta": {},  # No billed_units
        }
        mock_response.json = MagicMock(return_value=response_json)

        input_data = ["test text"]
        data = {"texts": input_data, "input_type": "search_query"}
        model_response = EmbeddingResponse()

        result = self.config._transform_response(
            response=mock_response,
            api_key="test-api-key",
            logging_obj=self.logging_obj,
            data=data,
            model_response=model_response,
            model=self.model,
            encoding=self.encoding,
            input=input_data,
        )

        # Verify usage uses encoding (mocked to return 5 tokens)
        assert result.usage is not None
        assert result.usage.prompt_tokens == 5
        assert result.usage.total_tokens == 5
        assert result.usage.completion_tokens == 0
        assert result.usage.prompt_tokens_details is None

        # Verify encoding was called
        self.encoding.encode.assert_called()
