import httpx
import json
import pytest
import sys
from typing import Any, Dict, List
from unittest.mock import MagicMock, Mock, patch
import os
from litellm._uuid import uuid
import time
import base64

sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system path
import litellm
from abc import ABC, abstractmethod
from litellm.integrations.custom_logger import CustomLogger
import json
from litellm.types.utils import StandardLoggingPayload


class BaseVectorStoreTest(ABC):
    """
    Abstract base test class that enforces a common test across all test classes.
    """

    @abstractmethod
    def get_base_request_args(self) -> dict:
        """Must return the base request args"""
        pass

    @abstractmethod
    def get_base_create_vector_store_args(self) -> dict:
        """Must return the base create vector store args"""
        pass

    @pytest.mark.parametrize("sync_mode", [True, False])
    @pytest.mark.asyncio
    async def test_basic_search_vector_store(self, sync_mode):
        litellm._turn_on_debug()
        litellm.set_verbose = True
        base_request_args = self.get_base_request_args()
        default_query = base_request_args.pop("query", "Basic ping")
        try:
            if sync_mode:
                response = litellm.vector_stores.search(
                    query=default_query, **base_request_args
                )
            else:
                response = await litellm.vector_stores.asearch(
                    query=default_query, **base_request_args
                )
        except litellm.InternalServerError:
            pytest.skip("Skipping test due to litellm.InternalServerError")

        print("litellm response=", json.dumps(response, indent=4, default=str))

        # Validate response structure
        self._validate_vector_store_response(response)

    @pytest.mark.parametrize("sync_mode", [True, False])
    @pytest.mark.asyncio
    async def test_basic_create_vector_store(self, sync_mode):
        litellm._turn_on_debug()
        litellm.set_verbose = True
        base_request_args = self.get_base_create_vector_store_args()

        # Extract custom_llm_provider from base args if present
        create_args = base_request_args
        try:
            if sync_mode:
                response = litellm.vector_stores.create(
                    name="Test Vector Store", **create_args
                )
            else:
                response = await litellm.vector_stores.acreate(
                    name="Test Vector Store", **create_args
                )
        except litellm.InternalServerError:
            pytest.skip("Skipping test due to litellm.InternalServerError")
        except Exception as e:
            # If this is an authentication or permission error, skip the test
            if (
                "authentication" in str(e).lower()
                or "permission" in str(e).lower()
                or "unauthorized" in str(e).lower()
            ):
                pytest.skip(
                    f"Skipping test due to authentication/permission error: {e}"
                )
            raise

        print("litellm create response=", json.dumps(response, indent=4, default=str))

        # Validate response structure
        self._validate_vector_store_create_response(response)

    def _validate_vector_store_response(self, response):
        """Validate the structure and content of a vector store search response"""

        # Check that response is a dictionary
        assert isinstance(
            response, dict
        ), f"Response should be a dict, got {type(response)}"

        # Check required top-level fields
        required_fields = ["object", "search_query", "data"]
        for field in required_fields:
            assert field in response, f"Missing required field '{field}' in response"

        # Validate object field
        assert (
            response["object"] == "vector_store.search_results.page"
        ), f"Expected object to be 'vector_store.search_results.page', got '{response['object']}'"

        # Validate search_query field
        assert isinstance(
            response["search_query"], str
        ), f"search_query should be a list, got {type(response['search_query'])}"
        assert len(response["search_query"]) > 0, "search_query should not be empty"
        assert all(
            isinstance(query, str) for query in response["search_query"]
        ), "All items in search_query should be strings"

        # Validate data field
        assert isinstance(
            response["data"], list
        ), f"data should be a list, got {type(response['data'])}"

        # Validate each result in data
        for i, result in enumerate(response["data"]):
            self._validate_search_result(result, i)

        print(
            f"✅ Response validation passed: Found {len(response['data'])} search results"
        )

    def _validate_vector_store_create_response(self, response):
        """Validate the structure and content of a vector store create response"""

        # Check that response is a dictionary
        assert isinstance(
            response, dict
        ), f"Response should be a dict, got {type(response)}"

        # Check required top-level fields for create response
        required_fields = ["id", "object", "created_at"]
        for field in required_fields:
            assert (
                field in response
            ), f"Missing required field '{field}' in create response"

        # Validate object field
        assert (
            response["object"] == "vector_store"
        ), f"Expected object to be 'vector_store', got '{response['object']}'"

        # Validate id field
        assert isinstance(
            response["id"], str
        ), f"id should be a string, got {type(response['id'])}"
        assert len(response["id"]) > 0, "id should not be empty"
        assert response["id"].startswith(
            "vs_"
        ), f"id should start with 'vs_', got '{response['id']}'"

        # Validate created_at field
        assert isinstance(
            response["created_at"], int
        ), f"created_at should be an integer, got {type(response['created_at'])}"
        assert response["created_at"] > 0, "created_at should be a positive timestamp"

        # Validate optional fields if present
        if "name" in response:
            assert isinstance(
                response["name"], str
            ), f"name should be a string, got {type(response['name'])}"

        if "bytes" in response:
            assert isinstance(
                response["bytes"], int
            ), f"bytes should be an integer, got {type(response['bytes'])}"
            assert response["bytes"] >= 0, "bytes should be non-negative"

        if "file_counts" in response:
            self._validate_file_counts(response["file_counts"])

        if "status" in response:
            valid_statuses = ["expired", "in_progress", "completed"]
            assert (
                response["status"] in valid_statuses
            ), f"status should be one of {valid_statuses}, got '{response['status']}'"

        if "expires_at" in response and response["expires_at"] is not None:
            assert isinstance(
                response["expires_at"], int
            ), f"expires_at should be an integer, got {type(response['expires_at'])}"

        if "last_active_at" in response and response["last_active_at"] is not None:
            assert isinstance(
                response["last_active_at"], int
            ), f"last_active_at should be an integer, got {type(response['last_active_at'])}"

        if "metadata" in response and response["metadata"] is not None:
            assert isinstance(
                response["metadata"], dict
            ), f"metadata should be a dict, got {type(response['metadata'])}"

        print(
            f"✅ Create response validation passed: Vector store '{response['id']}' created successfully"
        )

    def _validate_file_counts(self, file_counts):
        """Validate file_counts structure"""
        assert isinstance(
            file_counts, dict
        ), f"file_counts should be a dict, got {type(file_counts)}"

        required_count_fields = [
            "in_progress",
            "completed",
            "failed",
            "cancelled",
            "total",
        ]
        for field in required_count_fields:
            assert (
                field in file_counts
            ), f"Missing required field '{field}' in file_counts"
            assert isinstance(
                file_counts[field], int
            ), f"{field} should be an integer, got {type(file_counts[field])}"
            assert file_counts[field] >= 0, f"{field} should be non-negative"

        # Validate that total equals sum of other counts
        calculated_total = (
            file_counts["in_progress"]
            + file_counts["completed"]
            + file_counts["failed"]
            + file_counts["cancelled"]
        )
        assert (
            file_counts["total"] == calculated_total
        ), f"total should equal sum of other counts ({calculated_total}), got {file_counts['total']}"

    def _validate_search_result(self, result, index):
        """Validate an individual search result"""

        # Check that result is a dictionary
        assert isinstance(
            result, dict
        ), f"Result {index} should be a dict, got {type(result)}"

        # Check required fields in each result
        required_result_fields = [
            "file_id",
            "filename",
            "score",
            "attributes",
            "content",
        ]
        for field in required_result_fields:
            assert (
                field in result
            ), f"Missing required field '{field}' in result {index}"

        # Validate file_id
        assert isinstance(
            result["file_id"], str
        ), f"file_id should be a string, got {type(result['file_id'])} in result {index}"
        assert (
            len(result["file_id"]) > 0
        ), f"file_id should not be empty in result {index}"

        # Validate filename
        assert isinstance(
            result["filename"], str
        ), f"filename should be a string, got {type(result['filename'])} in result {index}"
        assert (
            len(result["filename"]) > 0
        ), f"filename should not be empty in result {index}"

        # Validate score
        assert isinstance(
            result["score"], (int, float)
        ), f"score should be a number, got {type(result['score'])} in result {index}"
        assert (
            0.0 <= result["score"] <= 1.0
        ), f"score should be between 0.0 and 1.0, got {result['score']} in result {index}"

        # Validate attributes
        assert isinstance(
            result["attributes"], dict
        ), f"attributes should be a dict, got {type(result['attributes'])} in result {index}"

        # Validate content
        assert isinstance(
            result["content"], list
        ), f"content should be a list, got {type(result['content'])} in result {index}"
        assert (
            len(result["content"]) > 0
        ), f"content should not be empty in result {index}"

        # Validate each content item
        for j, content_item in enumerate(result["content"]):
            assert isinstance(
                content_item, dict
            ), f"Content item {j} in result {index} should be a dict, got {type(content_item)}"
            assert (
                "type" in content_item
            ), f"Content item {j} in result {index} missing 'type' field"
            assert (
                "text" in content_item
            ), f"Content item {j} in result {index} missing 'text' field"
            assert isinstance(
                content_item["text"], str
            ), f"Content text should be a string in item {j} of result {index}"
            assert (
                len(content_item["text"]) > 0
            ), f"Content text should not be empty in item {j} of result {index}"

        print(
            f"✅ Result {index} validation passed: {result['filename']} (score: {result['score']:.4f})"
        )
