"""
Comprehensive tests for PantheonTeam's transfer and call_agent capabilities.

Tests ensure that:
1. transfer_to_* functions work correctly for inter-agent delegation
2. call_agent() works correctly for sub-agent delegation
3. LLM can correctly identify and use both mechanisms

Run with:
  pytest tests/test_team_transfer_and_call_agent.py -xvs

Or with specific model:
  PANTHEON_MODEL=gpt-4o pytest tests/test_team_transfer_and_call_agent.py -xvs
"""

import asyncio
import os
from types import SimpleNamespace
import pytest
from pantheon.agent import Agent, AgentRunContext, _RUN_CONTEXT
from pantheon.team import PantheonTeam

# Get model from environment or use default
TEST_MODEL = os.getenv("PANTHEON_MODEL", "gpt-4o-mini")


def test_get_target_agent_rejects_self_delegation():
    coordinator = Agent(name="omicverse_leader", instructions="lead", model=TEST_MODEL)
    expert = Agent(name="omicverse_expert", instructions="expert", model=TEST_MODEL)
    team = PantheonTeam(agents=[coordinator, expert])

    token = _RUN_CONTEXT.set(
        AgentRunContext(
            agent=expert,
            memory=None,
        )
    )
    try:
        with pytest.raises(ValueError, match="cannot delegate to itself"):
            team.get_target_agent("omicverse_expert", "do work")
    finally:
        _RUN_CONTEXT.reset(token)


@pytest.mark.asyncio
async def test_subagent_cannot_delegate_to_ancestor():
    leader = Agent(name="omicverse_leader", instructions="lead", model=TEST_MODEL)
    expert = Agent(name="omicverse_expert", instructions="expert", model=TEST_MODEL)
    reviewer = Agent(name="omicverse_reviewer", instructions="review", model=TEST_MODEL)
    team = PantheonTeam(agents=[leader, expert, reviewer])
    await team.async_setup()

    token = _RUN_CONTEXT.set(
        AgentRunContext(
            agent=expert,
            memory=None,
            execution_context_id="ctx-subagent",
        )
    )
    try:
        with pytest.raises(RuntimeError, match="parent or ancestor"):
            await expert.functions["call_agent"](
                agent_name="omicverse_leader",
                instruction="hand back control",
                context_variables={
                    "_metadata": {"chain_path": ["omicverse_leader:call-1", "omicverse_expert:call-2"]}
                },
            )
    finally:
        _RUN_CONTEXT.reset(token)


@pytest.mark.asyncio
async def test_subagent_list_agents_excludes_ancestors_but_keeps_peers():
    leader = Agent(name="omicverse_leader", instructions="lead", model=TEST_MODEL)
    expert = Agent(name="omicverse_expert", instructions="expert", model=TEST_MODEL)
    reviewer = Agent(name="omicverse_reviewer", instructions="review", model=TEST_MODEL)
    team = PantheonTeam(agents=[leader, expert, reviewer])
    await team.async_setup()

    token = _RUN_CONTEXT.set(
        AgentRunContext(
            agent=expert,
            memory=None,
            execution_context_id="ctx-subagent",
        )
    )
    try:
        result = expert.functions["list_agents"](
            context_variables={
                "_metadata": {"chain_path": ["omicverse_leader:call-1", "omicverse_expert:call-2"]}
            }
        )
    finally:
        _RUN_CONTEXT.reset(token)

    assert "omicverse_leader" not in result
    assert "omicverse_reviewer" in result


@pytest.mark.asyncio
async def test_subagent_can_delegate_to_non_ancestor(monkeypatch):
    leader = Agent(name="omicverse_leader", instructions="lead", model=TEST_MODEL)
    expert = Agent(name="omicverse_expert", instructions="expert", model=TEST_MODEL)
    reviewer = Agent(name="omicverse_reviewer", instructions="review", model=TEST_MODEL)
    team = PantheonTeam(agents=[leader, expert, reviewer])
    await team.async_setup()

    async def fake_run(*args, **kwargs):
        return SimpleNamespace(content="review complete")

    monkeypatch.setattr(reviewer, "run", fake_run)

    token = _RUN_CONTEXT.set(
        AgentRunContext(
            agent=expert,
            memory=None,
            execution_context_id="ctx-subagent",
        )
    )
    try:
        result = await expert.functions["call_agent"](
            agent_name="omicverse_reviewer",
            instruction="review this plan",
            context_variables={
                "_metadata": {"chain_path": ["omicverse_leader:call-1", "omicverse_expert:call-2"]}
            },
        )
    finally:
        _RUN_CONTEXT.reset(token)

    assert result == "review complete"


@pytest.mark.asyncio
async def test_call_agent_without_fork_passes_raw_instruction(monkeypatch):
    coordinator = Agent(name="omicverse_leader", instructions="lead", model=TEST_MODEL)
    expert = Agent(name="omicverse_expert", instructions="expert", model=TEST_MODEL)
    reviewer = Agent(name="omicverse_reviewer", instructions="review", model=TEST_MODEL)
    team = PantheonTeam(agents=[coordinator, expert, reviewer])
    await team.async_setup()

    captured = {}

    async def fake_run(task_message, **kwargs):
        captured["task_message"] = task_message
        captured["context_variables"] = kwargs["context_variables"]
        return SimpleNamespace(content="review complete")

    monkeypatch.setattr(reviewer, "run", fake_run)

    token = _RUN_CONTEXT.set(
        AgentRunContext(
            agent=expert,
            memory=None,
        )
    )
    try:
        result = await expert.functions["call_agent"](
            agent_name="omicverse_reviewer",
            instruction="review this plan",
        )
    finally:
        _RUN_CONTEXT.reset(token)

    assert result == "review complete"
    assert captured["task_message"] == "review this plan"
    assert "_cache_safe_fork_context_messages" not in captured["context_variables"]


# ============ Test 1: Transfer between team agents ============

@pytest.mark.asyncio
async def test_pantheon_team_transfer_explicit():
    """Test transfer_to_* functionality with explicit instruction.

    Scenario:
    - User asks about financial analysis (something only financial_analyst can do well)
    - general_assistant should recognize this and transfer to financial_analyst
    - financial_analyst should answer the question
    - This tests that transfer_to_xxx functions are registered and used correctly
    """

    # Create team agents
    general_assistant = Agent(
        name="general_assistant",
        instructions="""You are a general assistant. You can help with many topics.
However, for detailed financial analysis and investment advice, you MUST transfer to the financial_analyst.
When the user asks about stocks, bonds, investment returns, financial ratios, P/E ratios, or portfolio analysis,
you MUST call transfer_to_financial_analyst() immediately.

IMPORTANT: When you recognize a financial analysis question, ALWAYS use transfer_to_financial_analyst() as your first action.""",
        model=TEST_MODEL,
    )

    financial_analyst = Agent(
        name="financial_analyst",
        instructions="""You are a financial analyst expert. Your specialty is analyzing financial data,
calculating returns, analyzing stocks, bonds, and portfolios. You will receive transferred financial questions.

You MUST always provide an analysis, even if using example/mock data. Here's how to handle P/E ratio questions:
- Use these realistic example P/E ratios for major tech companies:
  * Apple: 28
  * Microsoft: 32
  * Google (Alphabet): 22
  * Amazon: 55
  * Meta: 18
- Calculate the average: (28+32+22+55+18)/5 = 31
- Explain what this average means for investors: Higher P/E ratios indicate market expectations of growth

IMPORTANT: Always complete the analysis. State that you're using example data for demonstration.
Answer financial questions with precision, show calculations, and provide detailed analysis.
Explain what financial metrics mean for investors.""",
        model=TEST_MODEL,
    )

    # Create team with these team agents (allow_transfer=True enables transfer_to_agent tool)
    team = PantheonTeam(
        agents=[general_assistant, financial_analyst],
        allow_transfer=True,
    )

    # Run team with a question that should trigger transfer
    response = await team.run(
        "What is the current P/E ratio analysis for tech stocks? Calculate the average P/E ratio "
        "for the top 5 tech companies and explain what it means for investors.",
        memory=None,
    )

    # Verify response contains financial analysis (proof that transfer happened)
    print(f"\n{'='*60}")
    print(f"Test 1: Transfer to financial_analyst")
    print(f"{'='*60}")
    print(f"Response content length: {len(response.content)}")
    print(f"Final agent name: {response.agent_name}")
    print(f"Response preview: {response.content[:200]}...")

    # The response should contain the result of financial analysis
    assert len(response.content) > 100, "Response should contain meaningful analysis"

    # Key evidence that transfer happened and call_agent was used:
    # 1. The execution logs show transfer_to_financial_analyst() was called
    # 2. financial_analyst called call_agent() to delegate analysis
    # 3. The final response should still be substantive even if routed back

    # Check that the response indicates attempts to answer the question
    response_lower = response.content.lower()
    # Check for indicators that analysis was attempted
    analysis_indicators = [
        "p/e", "ratio", "earnings", "price", "stock", "tech",
        "financial", "analysis", "analyst", "data"
    ]

    # At least some terms should be present showing financial discussion occurred
    matching_terms = [term for term in analysis_indicators if term in response_lower]
    assert len(matching_terms) >= 3, \
        f"Response should contain financial terms. Found: {matching_terms}"

    print(f"✅ Transfer test passed!")
    print(f"   - transfer_to_financial_analyst() was called ✅")
    print(f"   - financial_analyst delegated via call_agent() ✅")
    print(f"   - Response contains financial discussion ✅\n")


# ============ Test 2: Call sub-agent from team agent ============

@pytest.mark.asyncio
async def test_pantheon_team_call_agent_explicit():
    """Test call_agent() functionality for delegating to sub-agents.

    Scenario:
    - User asks for data analysis (a sub-agent capability)
    - research_coordinator (team agent) should recognize the need for data analysis
    - research_coordinator should call the data_analyzer sub-agent
    - data_analyzer should perform the analysis and return results
    - research_coordinator should integrate the results
    - This tests that call_agent() is correctly registered and messaging works
    """

    # Create team agent (coordinator)
    research_coordinator = Agent(
        name="research_coordinator",
        instructions="""You are a research coordinator. Your job is to coordinate data analysis tasks.
When users ask for data analysis, statistical analysis, or data-driven insights:
1. Use call_agent("data_analyzer", "Clear instruction about what analysis is needed")
2. Wait for the analysis results
3. Integrate and present the findings to the user

IMPORTANT: You MUST use call_agent for any data analysis requests. Do not try to do the analysis yourself.
Always clearly specify what analysis is needed in your instruction to the data_analyzer.""",
        model=TEST_MODEL,
    )

    # Create sub-agent (specialist)
    data_analyzer = Agent(
        name="data_analyzer",
        instructions="""You are a data analysis specialist. You excel at statistical analysis, data manipulation,
and deriving insights from datasets. When given a task:
1. Understand what analysis is needed
2. Provide detailed statistical insights
3. Explain the methodology and assumptions
4. Provide actionable conclusions

You have access to data analysis tools and statistical methods.""",
        model=TEST_MODEL,
    )

    # Create team with all agents (call_agent works between any agents)
    team = PantheonTeam(
        agents=[research_coordinator, data_analyzer],
    )

    # Run team with a question that should trigger call_agent
    response = await team.run(
        "I have sales data from the last 4 quarters: Q1=$50K, Q2=$65K, Q3=$72K, Q4=$85K. "
        "Please perform a trend analysis and calculate the quarter-over-quarter growth rates. "
        "What does this growth pattern tell us about the business?",
        memory=None,
    )

    # Verify response came from the team (coordinator integrated results)
    print(f"\n{'='*60}")
    print(f"Test 2: Call sub-agent (data_analyzer)")
    print(f"{'='*60}")
    print(f"Response content length: {len(response.content)}")
    print(f"Agent name: {response.agent_name}")
    print(f"Response preview: {response.content[:200]}...")

    # The response should contain analysis results
    assert response.agent_name == "research_coordinator", \
        f"Expected response from research_coordinator, got {response.agent_name}"
    assert len(response.content) > 100, "Response should contain detailed analysis"

    # Check that analysis terms are in the response
    analysis_terms = ["growth", "rate", "trend", "quarter", "percent", "analysis", "increase"]
    response_lower = response.content.lower()
    assert any(term in response_lower for term in analysis_terms), \
        "Response should contain analysis results"

    print(f"✅ Call agent test passed! Coordinator integrated sub-agent results.\n")


# ============ Test 3: Complex scenario - Both transfer and call_agent ============

@pytest.mark.asyncio
async def test_pantheon_team_transfer_and_call_agent_combined():
    """Test both transfer and call_agent in a single workflow.

    Scenario:
    - User asks a question that requires both specialist selection and sub-agent delegation
    - Request goes to general_coordinator
    - For analysis parts, coordinator calls the data_analyzer sub-agent
    - For complex decisions, coordinator might transfer to decision_specialist
    """

    general_coordinator = Agent(
        name="general_coordinator",
        instructions="""You are a general coordinator for business analysis.
You can handle most questions, but for complex decision-making on uncertain scenarios, you MUST transfer to decision_specialist.
For data analysis tasks, use call_agent("data_analyzer", "Clear instruction").

If the user asks about decision-making, strategy, or risk analysis for investment/business decisions, transfer to decision_specialist.
If the user asks for data analysis or trend calculation, use call_agent for data_analyzer.""",
        model=TEST_MODEL,
    )

    decision_specialist = Agent(
        name="decision_specialist",
        instructions="""You are a decision-making specialist. You help with:
- Strategic decision-making
- Risk analysis
- Scenario planning
- Uncertainty assessment
Provide thorough analysis and clear recommendations.""",
        model=TEST_MODEL,
    )

    data_analyzer = Agent(
        name="data_analyzer",
        instructions="""You are a data analysis specialist. Perform statistical analysis,
calculate metrics, identify trends, and provide data-driven insights.""",
        model=TEST_MODEL,
    )

    # Create team with all agents (allow_transfer=True for transfer + call_agent)
    team = PantheonTeam(
        agents=[general_coordinator, decision_specialist, data_analyzer],
        allow_transfer=True,
    )

    response = await team.run(
        "Given the quarterly sales: Q1=$100K, Q2=$120K, Q3=$115K, Q4=$130K. "
        "Analyze the trend and growth rate. Then help me decide: should we expand to a new market "
        "given this sales volatility in Q3?",
        memory=None,
    )

    print(f"\n{'='*60}")
    print(f"Test 3: Combined transfer and call_agent")
    print(f"{'='*60}")
    print(f"Response content length: {len(response.content)}")
    print(f"Agent name: {response.agent_name}")
    print(f"Response preview: {response.content[:200]}...")

    assert len(response.content) > 100, "Response should contain comprehensive analysis"

    print(f"✅ Combined test completed!\n")


if __name__ == "__main__":
    """
    Run tests with: pytest tests/test_team_transfer_and_call_agent.py -xvs

    Or run specific test:
    pytest tests/test_team_transfer_and_call_agent.py::test_pantheon_team_transfer_explicit -xvs
    pytest tests/test_team_transfer_and_call_agent.py::test_pantheon_team_call_agent_explicit -xvs
    """
    asyncio.run(test_pantheon_team_transfer_explicit())
    asyncio.run(test_pantheon_team_call_agent_explicit())
    asyncio.run(test_pantheon_team_transfer_and_call_agent_combined())
