"""E2E test for real AgentRunner + PydanticAI TestModel lifecycle.

Verifies:
1. State transitions: pending -> running -> completed
2. response_text extraction from agent result
3. Token usage captured across turns

Uses real AgentRunner with PydanticAI TestModel (no mocked runner).

Spec source:
- e2e-coverage.real-agent-lifecycle
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import pytest
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from returns.result import Success

from anima.core.lifecycle import LifecycleEvent
from anima.core.schemas import (
    AgentResult,
    AgentState,
    CompletionPolicy,
    JobSpec,
    PersonaSpec,
)
from anima.shell.agent import AgentRunner, create_agent
from anima.shell.state import create_agent_state, read_state, transition_state

# Import unified fake MCP server from helpers
from tests.helpers.mcp_fakes import FakeMCPServer


@pytest.fixture
async def test_agent(tmp_state_dir: Path) -> Agent[Any]:
    """Create a test agent with TestModel.

    Args:
        tmp_state_dir: Temporary state directory fixture.

    Returns:
        Configured PydanticAI Agent with TestModel.
    """
    _ = tmp_state_dir
    persona = PersonaSpec(
        id="test-persona", description="A test persona.",
        prompt="You are a test agent.",
        model="test-model",
    )

    result = await create_agent(
        persona,
        model_override="test",  # Force TestModel usage
    )

    assert isinstance(result, Success), f"create_agent failed: {result.failure()}"
    return result.unwrap()


@pytest.mark.asyncio
async def test_agent_lifecycle_state_transitions(
    tmp_state_dir: Path,
) -> None:
    """Real AgentRunner + TestModel: pending -> running -> completed transitions.

    Verifies:
    - Initial state is pending
    - After start, state is running
    - After completion, state is completed

    Args:
        tmp_state_dir: Temporary state directory fixture.
    """
    _ = tmp_state_dir
    job = JobSpec(
        persona_ref="test@abc123",
        instance_id="inst_lifecycle_test",
        max_turns=10,
    )

    # Create initial state
    created = create_agent_state("inst_lifecycle_test", "test-persona", job)
    assert isinstance(created, Success)
    state = created.unwrap()

    # Verify initial state is pending
    assert state.state == AgentState.pending
    assert state.instance_id == "inst_lifecycle_test"

    # Transition to running
    started = transition_state("inst_lifecycle_test", LifecycleEvent.start)
    assert isinstance(started, Success)

    # Verify running state
    running_result = read_state("inst_lifecycle_test")
    assert isinstance(running_result, Success)
    running_state = running_result.unwrap()
    assert running_state.state == AgentState.running

    # Transition to completed
    finished = transition_state("inst_lifecycle_test", LifecycleEvent.finish)
    assert isinstance(finished, Success)

    # Verify completed state
    completed_result = read_state("inst_lifecycle_test")
    assert isinstance(completed_result, Success)
    completed_state = completed_result.unwrap()
    assert completed_state.state == AgentState.completed


@pytest.mark.asyncio
async def test_agent_response_text_extraction(
    tmp_state_dir: Path,
) -> None:
    """Real AgentRunner + TestModel extracts response_text from agent result.

    Verifies:
    - AgentRunner.run() returns Success[AgentResult]
    - AgentResult contains response_text
    - response_text matches TestModel output

    Args:
        tmp_state_dir: Temporary state directory fixture.
    """
    _ = tmp_state_dir
    test_response = "This is a test response from the agent."

    persona = PersonaSpec(
        id="test-persona", description="A test persona.",
        prompt="You are a test agent.",
        model="test-model",
    )

    job = JobSpec(
        persona_ref="test@response",
        instance_id="inst_response_test",
        max_turns=5,
    )

    # Create agent state
    created = create_agent_state("inst_response_test", "test-persona", job)
    assert isinstance(created, Success)

    # Create RunnerContext with TestModel
    runner = AgentRunner(
        instance_id="inst_response_test",
        task="Simple test task",
        persona=persona,
        job=job,
        completion_policy=CompletionPolicy.auto,
        model_override="test",  # Force TestModel usage
    )

    # Create TestModel agent
    test_model = TestModel(custom_output_text=test_response)
    agent = Agent(
        model=test_model,
        system_prompt=persona.prompt,
    )

    # Initialize run state
    started = transition_state("inst_response_test", LifecycleEvent.start)
    assert isinstance(started, Success)

    # Run agent loop directly (bypassing full run() which starts MCP server)
    from returns.result import Failure

    state_result = read_state("inst_response_test")
    assert isinstance(state_result, Success)
    state_data = state_result.unwrap()

    # Inject fake MCP server

    runner.mcp_server = FakeMCPServer()  # type: ignore[assignment]

    # Run the agent loop
    result = await runner._run_agent_loop(agent, state_data)

    # Verify result
    assert isinstance(result, Success), (
        f"Agent loop failed: {result.failure() if isinstance(result, Failure) else 'unknown'}"
    )
    agent_result = result.unwrap()

    # Verify AgentResult contains response_text
    assert agent_result.response_text is not None
    assert test_response in agent_result.response_text


@pytest.mark.asyncio
async def test_agent_token_usage_captured(
    tmp_state_dir: Path,
) -> None:
    """Real AgentRunner + TestModel captures token usage across turns.

    Verifies:
    - AgentResult contains token_usage
    - TokenUsage has prompt_tokens and completion_tokens > 0
    - TokenUsage persists through agent run

    Args:
        tmp_state_dir: Temporary state directory fixture.
    """
    _ = tmp_state_dir
    test_response = "Token usage test response."

    persona = PersonaSpec(
        id="test-persona", description="A test persona.",
        prompt="You are a test agent.",
        model="test-model",
    )

    job = JobSpec(
        persona_ref="test@tokens",
        instance_id="inst_token_test",
        max_turns=5,
    )

    # Create agent state
    created = create_agent_state("inst_token_test", "test-persona", job)
    assert isinstance(created, Success)

    # Create RunnerContext with TestModel
    runner = AgentRunner(
        instance_id="inst_token_test",
        task="Token tracking test",
        persona=persona,
        job=job,
        completion_policy=CompletionPolicy.auto,
        model_override="test",
    )

    # Create TestModel agent
    test_model = TestModel(custom_output_text=test_response)
    agent = Agent(
        model=test_model,
        system_prompt=persona.prompt,
    )

    # Initialize run state
    started = transition_state("inst_token_test", LifecycleEvent.start)
    assert isinstance(started, Success)

    # Get state data
    state_result = read_state("inst_token_test")
    assert isinstance(state_result, Success)
    state_data = state_result.unwrap()

    # Inject fake MCP server

    runner.mcp_server = FakeMCPServer()  # type: ignore[assignment]

    # Run the agent loop
    result = await runner._run_agent_loop(agent, state_data)

    # Verify result
    from returns.result import Failure

    assert isinstance(result, Success), (
        f"Agent loop failed: {result.failure() if isinstance(result, Failure) else 'unknown'}"
    )
    agent_result = result.unwrap()

    # Verify token usage captured
    assert agent_result.token_usage is not None
    assert agent_result.token_usage.prompt_tokens > 0, "prompt_tokens should be > 0"
    assert agent_result.token_usage.completion_tokens > 0, (
        "completion_tokens should be > 0"
    )
    assert agent_result.token_usage.total_tokens > 0, "total_tokens should be > 0"

    # Verify token usage matches sum
    expected_total = (
        agent_result.token_usage.prompt_tokens
        + agent_result.token_usage.completion_tokens
    )
    assert agent_result.token_usage.total_tokens == expected_total, (
        f"total_tokens ({agent_result.token_usage.total_tokens}) != prompt ({agent_result.token_usage.prompt_tokens}) + completion ({agent_result.token_usage.completion_tokens})"
    )


@pytest.mark.asyncio
async def test_agent_full_lifecycle_integration(
    tmp_state_dir: Path,
) -> None:
    """Complete E2E integration: state transitions + response + tokens in one run.

    Integration test verifying all three requirements together:
    1. pending -> running -> completed state transitions
    2. response_text extracted from agent result
    3. token_usage captured

    Args:
        tmp_state_dir: Temporary state directory fixture.
    """
    _ = tmp_state_dir
    test_response = "Complete E2E integration test response."

    persona = PersonaSpec(
        id="test-persona", description="A test persona.",
        prompt="You are a helpful test assistant.",
        model="test-model",
    )

    job = JobSpec(
        persona_ref="test@integration",
        instance_id="inst_e2e_integration",
        max_turns=3,
    )

    # Create agent state
    created = create_agent_state("inst_e2e_integration", "test-persona", job)
    assert isinstance(created, Success)
    instance_id = created.unwrap().instance_id

    # Verify pending state
    pending_result = read_state(instance_id)
    assert isinstance(pending_result, Success)
    pending_state = pending_result.unwrap()
    assert pending_state.state == AgentState.pending

    # Create RunnerContext with TestModel
    runner = AgentRunner(
        instance_id=instance_id,
        task="Integration test task",
        persona=persona,
        job=job,
        completion_policy=CompletionPolicy.auto,
        model_override="test",
    )

    # Create TestModel agent
    test_model = TestModel(custom_output_text=test_response)
    agent = Agent(
        model=test_model,
        system_prompt=persona.prompt,
    )

    # Initialize run state
    started = transition_state(instance_id, LifecycleEvent.start)
    assert isinstance(started, Success)

    # Verify running state
    running_result = read_state(instance_id)
    assert isinstance(running_result, Success)
    running_state = running_result.unwrap()
    assert running_state.state == AgentState.running

    # Get state data for run loop
    state_result = read_state(instance_id)
    assert isinstance(state_result, Success)
    state_data = state_result.unwrap()

    # Inject fake MCP server

    runner.mcp_server = FakeMCPServer()  # type: ignore[assignment]

    # Run the agent loop
    from returns.result import Failure

    result = await runner._run_agent_loop(agent, state_data)

    # Step 1: Verify result success
    assert isinstance(result, Success), (
        f"Agent loop failed: {result.failure() if isinstance(result, Failure) else 'unknown'}"
    )
    agent_result = result.unwrap()
    assert isinstance(agent_result, AgentResult)

    # Step 2: Verify state transitioned to completed
    assert agent_result.final_state == AgentState.completed
    completed_result = read_state(instance_id)
    assert isinstance(completed_result, Success)
    completed_state = completed_result.unwrap()
    assert completed_state.state == AgentState.completed

    # Step 3: Verify response_text extracted
    assert agent_result.response_text is not None
    assert test_response in agent_result.response_text

    # Step 4: Verify token_usage captured
    assert agent_result.token_usage is not None
    assert agent_result.token_usage.prompt_tokens > 0
    assert agent_result.token_usage.completion_tokens > 0
    assert agent_result.token_usage.total_tokens > 0

    # Lifecycle assertion matrix:
    # - pending -> running: PASSED (started transition succeeded)
    # - running -> completed: PASSED (final_state == completed)
    # - response_text: PASSED (not None, contains test_response)
    # - token_usage: PASSED (all components > 0)
