"""
Test suite for OpenRouter API support in TradingAgents.

This module tests:
1. OpenRouter provider initialization with ChatOpenAI
2. API key handling (OPENROUTER_API_KEY vs OPENAI_API_KEY)
3. Error handling for missing API keys
4. Model name format validation (provider/model-name)
5. Embedding fallback behavior
6. Configuration validation
"""

import os
import pytest
from unittest.mock import Mock, patch, MagicMock
from typing import Dict, Any

# Import modules under test
from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.agents.utils.memory import FinancialSituationMemory
from tradingagents.default_config import DEFAULT_CONFIG


# ============================================================================
# Fixtures
# ============================================================================

@pytest.fixture
def openrouter_config():
    """Create a valid OpenRouter configuration."""
    config = DEFAULT_CONFIG.copy()
    config.update({
        "llm_provider": "openrouter",
        "deep_think_llm": "anthropic/claude-sonnet-4",
        "quick_think_llm": "anthropic/claude-haiku-3.5",
        "backend_url": "https://openrouter.ai/api/v1",
    })
    return config


@pytest.fixture
def mock_env_openrouter():
    """Mock environment with OPENROUTER_API_KEY set."""
    with patch.dict(os.environ, {
        "OPENROUTER_API_KEY": "sk-or-test-key-123",
    }, clear=True):
        yield


@pytest.fixture
def mock_env_openai():
    """Mock environment with only OPENAI_API_KEY set."""
    with patch.dict(os.environ, {
        "OPENAI_API_KEY": "sk-test-key-456",
    }, clear=True):
        yield


@pytest.fixture
def mock_env_empty():
    """Mock environment with no API keys."""
    with patch.dict(os.environ, {}, clear=True):
        yield


@pytest.fixture
def mock_langchain_classes():
    """Mock LangChain chat model classes."""
    with patch("tradingagents.graph.trading_graph.ChatOpenAI") as mock_openai, \
         patch("tradingagents.graph.trading_graph.ChatAnthropic") as mock_anthropic, \
         patch("tradingagents.graph.trading_graph.ChatGoogleGenerativeAI") as mock_google:

        # Configure mocks to return Mock instances
        mock_openai.return_value = Mock()
        mock_anthropic.return_value = Mock()
        mock_google.return_value = Mock()

        yield {
            "openai": mock_openai,
            "anthropic": mock_anthropic,
            "google": mock_google,
        }


@pytest.fixture
def mock_memory():
    """Mock FinancialSituationMemory to avoid actual ChromaDB/OpenAI calls."""
    with patch("tradingagents.graph.trading_graph.FinancialSituationMemory") as mock:
        mock.return_value = Mock(spec=FinancialSituationMemory)
        yield mock


@pytest.fixture
def mock_openai_client():
    """Mock OpenAI client for embedding tests."""
    with patch("tradingagents.agents.utils.memory.OpenAI") as mock:
        client_instance = Mock()
        mock.return_value = client_instance

        # Mock embedding response
        embedding_response = Mock()
        embedding_response.data = [Mock(embedding=[0.1] * 1536)]
        client_instance.embeddings.create.return_value = embedding_response

        yield mock


@pytest.fixture
def mock_chromadb():
    """Mock ChromaDB client."""
    with patch("tradingagents.agents.utils.memory.chromadb.Client") as mock:
        client_instance = Mock()
        collection_instance = Mock()
        collection_instance.count.return_value = 0
        client_instance.create_collection.return_value = collection_instance
        mock.return_value = client_instance
        yield mock


# ============================================================================
# Unit Tests: OpenRouter Provider Initialization
# ============================================================================

class TestOpenRouterInitialization:
    """Test OpenRouter provider initializes ChatOpenAI with correct parameters."""

    def test_openrouter_uses_chatopenai(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that OpenRouter provider uses ChatOpenAI class."""
        # Arrange: OpenRouter config ready

        # Act: Initialize TradingAgentsGraph
        graph = TradingAgentsGraph(config=openrouter_config)

        # Assert: ChatOpenAI was called, not Anthropic or Google
        assert mock_langchain_classes["openai"].called
        assert not mock_langchain_classes["anthropic"].called
        assert not mock_langchain_classes["google"].called

    def test_openrouter_sets_correct_base_url(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that OpenRouter sets base_url to https://openrouter.ai/api/v1."""
        # Arrange
        expected_url = "https://openrouter.ai/api/v1"

        # Act
        graph = TradingAgentsGraph(config=openrouter_config)

        # Assert: Check both deep and quick thinking LLMs
        calls = mock_langchain_classes["openai"].call_args_list
        assert len(calls) >= 2, "Expected at least 2 ChatOpenAI calls"

        # Check deep thinking LLM
        deep_call_kwargs = calls[0][1]  # Get keyword arguments
        assert deep_call_kwargs["base_url"] == expected_url

        # Check quick thinking LLM
        quick_call_kwargs = calls[1][1]
        assert quick_call_kwargs["base_url"] == expected_url

    def test_openrouter_uses_provider_model_format(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that OpenRouter accepts provider/model-name format."""
        # Arrange: Model names in provider/model format
        deep_model = "anthropic/claude-sonnet-4"
        quick_model = "anthropic/claude-haiku-3.5"

        # Act
        graph = TradingAgentsGraph(config=openrouter_config)

        # Assert: Model names passed correctly
        calls = mock_langchain_classes["openai"].call_args_list

        deep_call_kwargs = calls[0][1]
        assert deep_call_kwargs["model"] == deep_model

        quick_call_kwargs = calls[1][1]
        assert quick_call_kwargs["model"] == quick_model

    def test_openrouter_alternative_models(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test OpenRouter with different provider/model combinations."""
        # Arrange: Test various model formats
        test_models = [
            "google/gemini-2.0-flash-exp",
            "openai/gpt-4o",
            "meta-llama/llama-3.3-70b-instruct",
        ]

        for model in test_models:
            # Reset mocks
            mock_langchain_classes["openai"].reset_mock()

            # Update config
            config = openrouter_config.copy()
            config["deep_think_llm"] = model

            # Act
            graph = TradingAgentsGraph(config=config)

            # Assert
            call_kwargs = mock_langchain_classes["openai"].call_args_list[0][1]
            assert call_kwargs["model"] == model


# ============================================================================
# Unit Tests: API Key Handling
# ============================================================================

class TestAPIKeyHandling:
    """Test that OpenRouter uses OPENROUTER_API_KEY, not OPENAI_API_KEY."""

    def test_openrouter_requires_openrouter_api_key(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that OPENROUTER_API_KEY is available when using OpenRouter."""
        # Arrange: OPENROUTER_API_KEY is set

        # Act
        graph = TradingAgentsGraph(config=openrouter_config)

        # Assert: Can access the API key
        assert os.getenv("OPENROUTER_API_KEY") == "sk-or-test-key-123"

    def test_openai_api_key_not_used_for_openrouter(
        self,
        openrouter_config,
        mock_env_openai,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that OPENAI_API_KEY alone is not sufficient for openrouter provider."""
        # Arrange: Only OPENAI_API_KEY is set (not OPENROUTER_API_KEY)

        # Act & Assert: Should raise ValueError requiring OPENROUTER_API_KEY
        with pytest.raises(ValueError, match="OPENROUTER_API_KEY"):
            graph = TradingAgentsGraph(config=openrouter_config)

    def test_missing_api_key_raises_error(
        self,
        openrouter_config,
        mock_env_empty,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that missing OPENROUTER_API_KEY raises clear error."""
        # Arrange: No API keys in environment

        # Act & Assert: Should raise ValueError
        with pytest.raises(ValueError, match="OPENROUTER_API_KEY"):
            graph = TradingAgentsGraph(config=openrouter_config)


# ============================================================================
# Unit Tests: Error Handling
# ============================================================================

class TestErrorHandling:
    """Test error handling for missing API keys and invalid configurations."""

    def test_invalid_provider_raises_error(
        self,
        openrouter_config,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that invalid provider raises ValueError."""
        # Arrange: Invalid provider
        config = openrouter_config.copy()
        config["llm_provider"] = "invalid_provider"

        # Act & Assert
        with pytest.raises(ValueError, match="Unsupported LLM provider"):
            graph = TradingAgentsGraph(config=config)

    def test_empty_backend_url_handled(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test behavior with empty backend_url."""
        # Arrange
        config = openrouter_config.copy()
        config["backend_url"] = ""

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert: Empty string passed to ChatOpenAI
        call_kwargs = mock_langchain_classes["openai"].call_args_list[0][1]
        assert call_kwargs["base_url"] == ""

    def test_none_backend_url_handled(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test behavior with None backend_url."""
        # Arrange
        config = openrouter_config.copy()
        config["backend_url"] = None

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert: None passed to ChatOpenAI
        call_kwargs = mock_langchain_classes["openai"].call_args_list[0][1]
        assert call_kwargs["base_url"] is None


# ============================================================================
# Integration Tests: Model Format Validation
# ============================================================================

class TestModelFormatValidation:
    """Test that OpenRouter model names work correctly."""

    def test_anthropic_model_format(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test Anthropic models via OpenRouter (anthropic/*)."""
        # Arrange
        config = openrouter_config.copy()
        config["deep_think_llm"] = "anthropic/claude-opus-4"
        config["quick_think_llm"] = "anthropic/claude-sonnet-3.5"

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert
        calls = mock_langchain_classes["openai"].call_args_list
        assert calls[0][1]["model"] == "anthropic/claude-opus-4"
        assert calls[1][1]["model"] == "anthropic/claude-sonnet-3.5"

    def test_openai_model_format(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test OpenAI models via OpenRouter (openai/*)."""
        # Arrange
        config = openrouter_config.copy()
        config["deep_think_llm"] = "openai/gpt-4o"
        config["quick_think_llm"] = "openai/gpt-4o-mini"

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert
        calls = mock_langchain_classes["openai"].call_args_list
        assert calls[0][1]["model"] == "openai/gpt-4o"
        assert calls[1][1]["model"] == "openai/gpt-4o-mini"

    def test_google_model_format(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test Google models via OpenRouter (google/*)."""
        # Arrange
        config = openrouter_config.copy()
        config["deep_think_llm"] = "google/gemini-2.0-flash-exp"
        config["quick_think_llm"] = "google/gemini-flash-1.5"

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert
        calls = mock_langchain_classes["openai"].call_args_list
        assert calls[0][1]["model"] == "google/gemini-2.0-flash-exp"
        assert calls[1][1]["model"] == "google/gemini-flash-1.5"

    def test_meta_llama_model_format(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test Meta Llama models via OpenRouter (meta-llama/*)."""
        # Arrange
        config = openrouter_config.copy()
        config["deep_think_llm"] = "meta-llama/llama-3.3-70b-instruct"
        config["quick_think_llm"] = "meta-llama/llama-3.1-8b-instruct"

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert
        calls = mock_langchain_classes["openai"].call_args_list
        assert calls[0][1]["model"] == "meta-llama/llama-3.3-70b-instruct"
        assert calls[1][1]["model"] == "meta-llama/llama-3.1-8b-instruct"


# ============================================================================
# Integration Tests: Embedding Handling
# ============================================================================

class TestEmbeddingHandling:
    """Test that embeddings work correctly with OpenRouter."""

    def test_memory_uses_openrouter_base_url(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_openai_client,
        mock_chromadb
    ):
        """Test that FinancialSituationMemory uses fallback OpenAI for embeddings."""
        # Arrange - need OPENAI_API_KEY for embeddings when using OpenRouter
        with patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-openai-key"}):
            config = openrouter_config.copy()

            # Act
            memory = FinancialSituationMemory("test_memory", config)

            # Assert: OpenAI client initialized (for embeddings fallback)
            assert mock_openai_client.called

    def test_memory_embedding_with_openrouter(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_openai_client,
        mock_chromadb
    ):
        """Test that embeddings can be generated via OpenAI fallback."""
        # Arrange - need OPENAI_API_KEY for embeddings
        with patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-openai-key"}):
            memory = FinancialSituationMemory("test_memory", openrouter_config)
            test_text = "Test financial situation"

            # Act
            embedding = memory.get_embedding(test_text)

            # Assert: Embedding created successfully
            assert embedding is not None
            assert len(embedding) == 1536
            mock_openai_client.return_value.embeddings.create.assert_called_once()

    def test_memory_uses_text_embedding_model(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_openai_client,
        mock_chromadb
    ):
        """Test that correct embedding model is used."""
        # Arrange - need OPENAI_API_KEY for embeddings
        with patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-openai-key"}):
            memory = FinancialSituationMemory("test_memory", openrouter_config)

            # Act
            memory.get_embedding("test")

            # Assert: Called with text-embedding-3-small
            call_args = mock_openai_client.return_value.embeddings.create.call_args
            assert call_args[1]["model"] == "text-embedding-3-small"

    def test_memory_ollama_embedding_model(
        self,
        openrouter_config,
        mock_openai_client,
        mock_chromadb
    ):
        """Test that Ollama uses nomic-embed-text model."""
        # Arrange
        config = openrouter_config.copy()
        config["backend_url"] = "http://localhost:11434/v1"
        memory = FinancialSituationMemory("test_memory", config)

        # Act
        memory.get_embedding("test")

        # Assert: Called with nomic-embed-text
        call_args = mock_openai_client.return_value.embeddings.create.call_args
        assert call_args[1]["model"] == "nomic-embed-text"

    def test_memory_graceful_fallback_on_embedding_error(
        self,
        openrouter_config,
        mock_openai_client,
        mock_chromadb
    ):
        """Test that memory gracefully handles embedding failures."""
        # Arrange
        memory = FinancialSituationMemory("test_memory", openrouter_config)

        # Mock embedding failure
        mock_openai_client.return_value.embeddings.create.side_effect = Exception(
            "API quota exceeded"
        )

        # Act: Try to get memories (will fail on embedding)
        result = memory.get_memories("current situation", n_matches=1)

        # Assert: Returns empty list instead of crashing
        assert result == []

    def test_memory_add_situations_with_openrouter(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_openai_client,
        mock_chromadb
    ):
        """Test adding situations to memory using OpenRouter embeddings."""
        # Arrange - need OPENAI_API_KEY for embeddings
        with patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-openai-key"}):
            memory = FinancialSituationMemory("test_memory", openrouter_config)
            situations = [
                ("Market volatility increasing", "Reduce risk exposure"),
                ("Strong uptrend detected", "Increase position size"),
            ]

            # Act
            memory.add_situations(situations)

            # Assert: Embeddings created for each situation
            assert mock_openai_client.return_value.embeddings.create.call_count == 2


# ============================================================================
# Edge Cases
# ============================================================================

class TestEdgeCases:
    """Test edge cases and boundary conditions."""

    def test_case_insensitive_provider_name(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that provider names are case-insensitive.

        Provider names use .lower() comparison, so 'OpenRouter', 'OPENROUTER', etc. all work.
        """
        # Arrange: All cases should work
        valid_providers = ["openrouter", "OpenRouter", "OPENROUTER", "OpenRouTer"]

        for provider in valid_providers:
            # Reset mocks
            mock_langchain_classes["openai"].reset_mock()

            # Act
            config = openrouter_config.copy()
            config["llm_provider"] = provider
            graph = TradingAgentsGraph(config=config)

            # Assert: ChatOpenAI was called
            assert mock_langchain_classes["openai"].called

    def test_openrouter_with_ollama_provider_name(
        self,
        openrouter_config,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that 'ollama' provider also uses ChatOpenAI (grouped with openrouter)."""
        # Arrange
        config = openrouter_config.copy()
        config["llm_provider"] = "ollama"
        config["backend_url"] = "http://localhost:11434/v1"

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert: ChatOpenAI used (not Anthropic/Google)
        assert mock_langchain_classes["openai"].called
        assert not mock_langchain_classes["anthropic"].called
        assert not mock_langchain_classes["google"].called

    def test_empty_model_name(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test behavior with empty model names."""
        # Arrange
        config = openrouter_config.copy()
        config["deep_think_llm"] = ""
        config["quick_think_llm"] = ""

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert: Empty strings passed to ChatOpenAI
        calls = mock_langchain_classes["openai"].call_args_list
        assert calls[0][1]["model"] == ""
        assert calls[1][1]["model"] == ""

    def test_special_characters_in_model_name(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test model names with special characters."""
        # Arrange
        config = openrouter_config.copy()
        config["deep_think_llm"] = "anthropic/claude-3.5-sonnet-20241022"

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert: Model name preserved exactly
        call_kwargs = mock_langchain_classes["openai"].call_args_list[0][1]
        assert call_kwargs["model"] == "anthropic/claude-3.5-sonnet-20241022"

    def test_url_with_trailing_slash(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test backend_url with trailing slash."""
        # Arrange
        config = openrouter_config.copy()
        config["backend_url"] = "https://openrouter.ai/api/v1/"

        # Act
        graph = TradingAgentsGraph(config=config)

        # Assert: Trailing slash preserved
        call_kwargs = mock_langchain_classes["openai"].call_args_list[0][1]
        assert call_kwargs["base_url"] == "https://openrouter.ai/api/v1/"

    def test_memory_empty_collection_query(
        self,
        openrouter_config,
        mock_openai_client,
        mock_chromadb
    ):
        """Test querying memories when collection is empty."""
        # Arrange
        memory = FinancialSituationMemory("test_memory", openrouter_config)

        # Act: Query empty collection
        result = memory.get_memories("test situation", n_matches=5)

        # Assert: Returns empty list
        assert result == []

    def test_memory_zero_matches_requested(
        self,
        openrouter_config,
        mock_openai_client,
        mock_chromadb
    ):
        """Test requesting zero matches from memory."""
        # Arrange
        memory = FinancialSituationMemory("test_memory", openrouter_config)
        collection_mock = mock_chromadb.return_value.create_collection.return_value
        collection_mock.count.return_value = 5  # Non-empty collection
        collection_mock.query.return_value = {
            "documents": [[]],
            "metadatas": [[]],
            "distances": [[]]
        }

        # Act
        result = memory.get_memories("test", n_matches=0)

        # Assert: Returns empty list
        assert result == []


# ============================================================================
# Configuration Tests
# ============================================================================

class TestConfiguration:
    """Test configuration handling for OpenRouter."""

    def test_default_config_not_openrouter(self):
        """Test that default config doesn't use OpenRouter."""
        # Arrange & Act
        config = DEFAULT_CONFIG

        # Assert
        assert config["llm_provider"] != "openrouter"
        assert config["backend_url"] == "https://api.openai.com/v1"

    def test_config_override_with_openrouter(
        self,
        openrouter_config,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that config can be overridden to use OpenRouter."""
        # Arrange: OpenRouter config overrides defaults

        # Act
        graph = TradingAgentsGraph(config=openrouter_config)

        # Assert: Configuration applied
        assert graph.config["llm_provider"] == "openrouter"
        assert graph.config["backend_url"] == "https://openrouter.ai/api/v1"

    def test_partial_config_merge(
        self,
        mock_env_openrouter,
        mock_langchain_classes,
        mock_memory
    ):
        """Test that partial config requires all necessary keys.

        NOTE: Current implementation doesn't merge with defaults.
        Missing required keys like 'project_dir' will cause KeyError.
        User must provide complete config or use DEFAULT_CONFIG.copy().
        """
        # Arrange: Partial config missing required keys
        partial_config = {
            "llm_provider": "openrouter",
            "backend_url": "https://openrouter.ai/api/v1",
        }

        # Act & Assert: Missing 'project_dir' causes KeyError
        with pytest.raises(KeyError, match="project_dir"):
            graph = TradingAgentsGraph(config=partial_config)

        # Arrange: Complete config works
        complete_config = DEFAULT_CONFIG.copy()
        complete_config.update({
            "llm_provider": "openrouter",
            "backend_url": "https://openrouter.ai/api/v1",
        })

        # Act
        graph = TradingAgentsGraph(config=complete_config)

        # Assert: Uses provided values
        assert graph.config["llm_provider"] == "openrouter"
        assert graph.config["backend_url"] == "https://openrouter.ai/api/v1"