"""
Tests for the internal-agentic example.

Validates that the agentic AI example correctly demonstrates:
- Goal tracking with metrics
- Tool execution monitoring
- Human intervention tracking
- Reasoning step visibility
- Error handling and resilience
"""

import pytest
import sys
import os
from unittest.mock import patch, MagicMock
from pathlib import Path

# Add sdk-examples/internal-agentic to path
# Note: test is now in tests/integration/, so we need to go up two levels
example_dir = Path(__file__).parent.parent.parent / "sdk-examples" / "internal-agentic"
sys.path.insert(0, str(example_dir))

from agent_instrumented import MedicalResearchAgent
from tools import MockPubMedSearch, MockDrugDatabase, MockCalculator


@pytest.fixture
def mock_mca_client():
    """Create a mock MCAClient for testing."""
    client = MagicMock()
    client.record_goal_started = MagicMock(return_value="test-goal-123")
    client.record_goal_completed = MagicMock()
    client.track_tool = MagicMock()
    client.record_human_intervention = MagicMock()
    client.reasoning_step = MagicMock()
    client.shutdown = MagicMock()

    # Mock context manager for track_tool
    track_tool_context = MagicMock()
    track_tool_context.__enter__ = MagicMock(return_value=MagicMock())
    track_tool_context.__exit__ = MagicMock(return_value=False)
    client.track_tool.return_value = track_tool_context

    # Mock context manager for reasoning_step
    reasoning_step_context = MagicMock()
    reasoning_step_context.__enter__ = MagicMock(return_value=MagicMock())
    reasoning_step_context.__exit__ = MagicMock(return_value=False)
    client.reasoning_step.return_value = reasoning_step_context

    return client


class TestAgenticExample:
    """Test suite for agentic AI example."""

    def test_agent_initialization(self, mock_mca_client):
        """Test that agent initializes with tools."""
        agent = MedicalResearchAgent(mock_mca_client)

        assert agent.client == mock_mca_client
        assert isinstance(agent.pubmed, MockPubMedSearch)
        assert isinstance(agent.drug_db, MockDrugDatabase)
        assert isinstance(agent.calculator, MockCalculator)

    def test_goal_tracking_success_scenario(self, mock_mca_client):
        """Test that goal tracking metrics are recorded for successful goal."""
        agent = MedicalResearchAgent(mock_mca_client)

        question = "What are the latest treatments for Type 2 Diabetes?"
        result = agent.research_question(question)

        # Verify goal started was called
        mock_mca_client.record_goal_started.assert_called_once()
        call_args = mock_mca_client.record_goal_started.call_args
        assert call_args[1]["goal_description"] == question
        assert call_args[1]["goal_type"] == "research_question"

        # Verify goal completed with success
        mock_mca_client.record_goal_completed.assert_called_once_with(
            "test-goal-123", status="success"
        )

        # Verify result structure
        assert result["status"] == "success"
        assert result["goal_id"] == "test-goal-123"
        assert "answer" in result

    def test_goal_tracking_failure_scenario(self, mock_mca_client):
        """Test that goal tracking handles failures correctly.

        NOTE: PubMed search has fallback behavior, so we fail at drug database query
        which doesn't have fallback (critical failure point).
        """
        agent = MedicalResearchAgent(mock_mca_client, simulation_delay=0)

        # Make PubMed return results to trigger drug DB query, then fail at drug DB
        agent.pubmed.search = MagicMock(
            return_value=[{"title": "SGLT2 Test", "authors": "Test", "year": 2024}]
        )
        agent.drug_db.query = MagicMock(side_effect=Exception("Drug DB error"))

        question = "Test question"
        result = agent.research_question(question)

        # Verify goal completed with failure (drug DB is critical, no fallback)
        mock_mca_client.record_goal_completed.assert_called_once_with(
            "test-goal-123", status="failure"
        )

        # Verify result structure
        assert result["status"] == "failure"
        assert "error" in result

    def test_tool_execution_tracking(self, mock_mca_client):
        """Test that all tool calls are tracked."""
        agent = MedicalResearchAgent(mock_mca_client)

        question = "What are treatments for diabetes with cardiovascular complications?"
        agent.research_question(question)

        # Count tool calls (should be PubMed searches + drug database queries)
        tool_calls = mock_mca_client.track_tool.call_count
        assert tool_calls >= 2, "Expected at least 2 tool calls (PubMed + drug DB)"

        # Verify tool calls include goal_id
        for call in mock_mca_client.track_tool.call_args_list:
            assert "goal_id" in call[1], "Tool call should include goal_id"

    def test_human_intervention_tracking(self, mock_mca_client):
        """Test that human interventions are tracked."""
        agent = MedicalResearchAgent(mock_mca_client)

        question = "What are the latest treatments for Type 2 Diabetes?"
        agent.research_question(question)

        # Verify human intervention was recorded
        mock_mca_client.record_human_intervention.assert_called_once()
        call_args = mock_mca_client.record_human_intervention.call_args
        assert call_args[1]["reason"] == "clinical_decision_support"
        assert call_args[1]["intervention_type"] == "review_request"

    def test_reasoning_steps_visibility(self, mock_mca_client):
        """Test that reasoning steps are tracked with correct attributes."""
        agent = MedicalResearchAgent(mock_mca_client)

        question = "What are the latest treatments for Type 2 Diabetes?"
        agent.research_question(question)

        # Verify reasoning_step was called at least 3 times (planning, evaluation, synthesis)
        assert mock_mca_client.reasoning_step.call_count >= 3

        # Verify specific steps were called in order (checking call args)
        calls = mock_mca_client.reasoning_step.call_args_list

        # Check for planning step
        planning_call = next((c for c in calls if c[1].get("step_type") == "planning"), None)
        assert planning_call is not None, "Missing planning step"

        # Check for evaluation step
        evaluation_call = next((c for c in calls if c[1].get("step_type") == "evaluation"), None)
        assert evaluation_call is not None, "Missing evaluation step"

        # Check for synthesis step
        synthesis_call = next((c for c in calls if c[1].get("step_type") == "synthesis"), None)
        assert synthesis_call is not None, "Missing synthesis step"

    def test_multiple_agent_runs_independence(self, mock_mca_client):
        """Test that multiple agent runs work independently."""
        agent = MedicalResearchAgent(mock_mca_client)

        # Run twice with different goal IDs
        mock_mca_client.record_goal_started.side_effect = ["goal-1", "goal-2"]

        result1 = agent.research_question("Question 1")
        result2 = agent.research_question("Question 2")

        # Each run should have unique goal ID
        assert result1["goal_id"] != result2["goal_id"]

        # Both should complete
        assert result1["status"] == "success"
        assert result2["status"] == "success"

    def test_error_handling_tool_failure(self, mock_mca_client):
        """Test that agent handles tool failures gracefully."""
        agent = MedicalResearchAgent(mock_mca_client)

        # Make PubMed search return papers that will trigger drug DB query,
        # then make drug database fail
        agent.pubmed.search = MagicMock(
            return_value=[
                {
                    "title": "SGLT2 Inhibitors for Diabetes",
                    "authors": "Test",
                    "year": 2024,
                }
            ]
        )
        agent.drug_db.query = MagicMock(side_effect=Exception("Database timeout"))

        question = "Test question about diabetes"
        result = agent.research_question(question)

        # Agent should fail gracefully
        assert result["status"] == "failure"
        assert "error" in result
        assert "Database timeout" in result["error"]

    def test_literature_search_fallback_behavior(self, mock_mca_client):
        """Test that agent continues when one search term fails (fallback behavior)."""
        agent = MedicalResearchAgent(mock_mca_client, simulation_delay=0)

        # Make first search fail, second search succeed
        # Question about diabetes generates 2 search terms: "diabetes type 2 cardiovascular" and "treatment guidelines"
        search_results = [
            Exception("Network timeout"),  # First term fails
            [{"title": "SGLT2 Paper", "authors": "Author", "year": 2024}],  # Second succeeds
        ]
        agent.pubmed.search = MagicMock(side_effect=search_results)

        question = "What are treatments for diabetes with cardiovascular complications?"
        result = agent.research_question(question)

        # Agent should succeed despite first search failing (fallback behavior)
        assert result["status"] == "success"
        # Verify both search terms were attempted (fallback didn't stop execution)
        assert agent.pubmed.search.call_count >= 1

    def test_reasoning_step_exception_handling(self, mock_mca_client):
        """Test that reasoning_step() properly handles exceptions and records error status."""
        # Create a mock span that will be yielded by reasoning_step
        mock_span = MagicMock()
        mock_reasoning_step = MagicMock()
        mock_reasoning_step.__enter__ = MagicMock(return_value=mock_span)
        mock_reasoning_step.__exit__ = MagicMock(return_value=False)  # Don't suppress exception

        # Configure reasoning_step to yield the mock span
        mock_mca_client.reasoning_step.return_value = mock_reasoning_step

        # Create agent with simulation_delay=0 for fast testing
        agent = MedicalResearchAgent(mock_mca_client, simulation_delay=0)

        # Make ALL PubMed searches fail, then make drug DB fail too
        # This ensures the goal actually fails (can't rely on just PubMed failure due to fallback)
        agent.pubmed.search = MagicMock(side_effect=Exception("Search service unavailable"))
        agent.drug_db.query = MagicMock(side_effect=Exception("DB unavailable"))

        question = "Test question about diabetes"
        result = agent.research_question(question)

        # Even with search failures, agent may succeed with empty results due to fallback
        # The test verifies reasoning_step was called regardless of outcome
        assert mock_mca_client.reasoning_step.call_count >= 1

        # Verify the context manager was properly entered and exited
        assert mock_reasoning_step.__enter__.called
        assert mock_reasoning_step.__exit__.called


class TestMockTools:
    """Test suite for mock tools."""

    def test_pubmed_search(self):
        """Test PubMed mock search returns results."""
        pubmed = MockPubMedSearch()

        results = pubmed.search("diabetes type 2 cardiovascular")
        assert len(results) > 0
        assert all("title" in paper for paper in results)
        assert all("authors" in paper for paper in results)

    def test_drug_database_query(self):
        """Test drug database returns drug information."""
        drug_db = MockDrugDatabase()

        info = drug_db.query("sglt2_inhibitors")
        assert "class" in info
        assert "examples" in info
        assert "contraindications" in info
        assert len(info["examples"]) > 0

    def test_calculator(self):
        """Test calculator performs safe calculations."""
        calc = MockCalculator()

        result = calc.calculate("2 + 2")
        assert result == 4

        result = calc.calculate("10 * 5")
        assert result == 50


class TestAgenticMetricsIntegration:
    """Integration tests for agentic metrics (requires real MCAClient)."""

    @pytest.mark.integration
    @pytest.mark.skipif(
        not os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"),
        reason="Requires OTel collector - set OTEL_EXPORTER_OTLP_ENDPOINT to run",
    )
    def test_real_client_goal_metrics(self):
        """Test goal metrics with real MCAClient (integration test)."""
        # This test requires OpenTelemetry collector running
        # Run with: OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 pytest
        pass  # Placeholder for actual integration test

    @pytest.mark.integration
    @pytest.mark.skipif(
        not os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"),
        reason="Requires OTel collector - set OTEL_EXPORTER_OTLP_ENDPOINT to run",
    )
    def test_real_client_tool_metrics(self):
        """Test tool metrics with real MCAClient (integration test)."""
        # This test requires OpenTelemetry collector running
        # Run with: OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 pytest
        pass  # Placeholder for actual integration test
