//! Integration tests for `PromptResponse.messages` using mock models.
//! Exercises the real agent loop code path with mocked LLM responses.

use rig::agent::AgentBuilder;
use rig::completion::{Chat, Message, Prompt, Usage};
use rig::message::{AssistantContent, UserContent};
use rig::test_utils::{MockAddTool, MockCompletionModel, MockTurn};

// ---------------------------------------------------------------------------
// Mock model infrastructure
// ---------------------------------------------------------------------------

fn simple_text_turn() -> MockTurn {
    MockTurn::text("hello from mock")
        .with_usage(Usage {
            input_tokens: 10,
            output_tokens: 5,
            total_tokens: 15,
            cached_input_tokens: 0,
            cache_creation_input_tokens: 0,
            tool_use_prompt_tokens: 0,
            reasoning_tokens: 0,
        })
        .with_message_id("msg_mock_1")
}

fn simple_text_model(turns: usize) -> MockCompletionModel {
    MockCompletionModel::new((0..turns).map(|_| simple_text_turn()))
}

fn tool_then_text_model() -> MockCompletionModel {
    MockCompletionModel::new([
        MockTurn::tool_call("tc_1", "add", serde_json::json!({"x": 2, "y": 3}))
            .with_usage(Usage {
                input_tokens: 15,
                output_tokens: 8,
                total_tokens: 23,
                cached_input_tokens: 0,
                cache_creation_input_tokens: 0,
                tool_use_prompt_tokens: 0,
                reasoning_tokens: 0,
            })
            .with_message_id("msg_tool"),
        MockTurn::text("The answer is 5")
            .with_usage(Usage {
                input_tokens: 20,
                output_tokens: 4,
                total_tokens: 24,
                cached_input_tokens: 0,
                cache_creation_input_tokens: 0,
                tool_use_prompt_tokens: 0,
                reasoning_tokens: 0,
            })
            .with_message_id("msg_text"),
    ])
}

fn always_tool_call_turn() -> MockTurn {
    MockTurn::tool_call("tc_loop", "add", serde_json::json!({"x": 1, "y": 1}))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

/// Test 1: Standard path still returns a plain String (backward compat).
#[tokio::test]
async fn standard_prompt_returns_string() {
    let agent = AgentBuilder::new(simple_text_model(1)).build();

    let result: String = agent.prompt("hi").await.expect("prompt should succeed");

    assert_eq!(result, "hello from mock");
}

/// Test 2: `extended_details()` returns a `PromptResponse` with `messages: Some(...)`.
#[tokio::test]
async fn extended_details_populates_messages() {
    let agent = AgentBuilder::new(simple_text_model(1)).build();

    let resp = agent
        .prompt("hi")
        .extended_details()
        .await
        .expect("prompt should succeed");

    assert_eq!(resp.output, "hello from mock");
    assert_eq!(resp.usage.input_tokens, 10);
    assert_eq!(resp.usage.output_tokens, 5);

    // Messages should be populated
    let messages = resp
        .messages
        .expect("messages should be Some for extended_details");

    // Should contain: [User("hi"), Assistant("hello from mock")]
    assert_eq!(messages.len(), 2);

    // First message: User
    match &messages[0] {
        Message::User { content } => match content.first() {
            UserContent::Text(t) => assert_eq!(t.text, "hi"),
            other => panic!("expected text user content, got: {other:?}"),
        },
        other => panic!("expected User message, got: {other:?}"),
    }

    // Second message: Assistant
    match &messages[1] {
        Message::Assistant { content, .. } => match content.first() {
            AssistantContent::Text(t) => assert_eq!(t.text, "hello from mock"),
            other => panic!("expected text assistant content, got: {other:?}"),
        },
        other => panic!("expected Assistant message, got: {other:?}"),
    }
}

/// Test 3: `with_history()` + `extended_details()` — the response messages
/// should contain the full conversation including any provided history.
#[tokio::test]
async fn extended_with_history_both_populated() {
    let agent = AgentBuilder::new(simple_text_model(1)).build();

    let initial_history: Vec<Message> = Vec::new();

    let resp = agent
        .prompt("hello")
        .with_history(&initial_history)
        .extended_details()
        .await
        .expect("prompt should succeed");

    let response_messages = resp.messages.expect("messages should be Some");

    // Response should contain the full conversation (User + Assistant)
    assert_eq!(response_messages.len(), 2);

    // First message: User
    match &response_messages[0] {
        Message::User { .. } => {}
        other => panic!("expected User, got: {other:?}"),
    }

    // Second message: Assistant
    match &response_messages[1] {
        Message::Assistant { .. } => {}
        other => panic!("expected Assistant, got: {other:?}"),
    }
}

/// Test 4: Standard path with `with_history()` — verify the API works with
/// an immutable history reference.
#[tokio::test]
async fn standard_with_history_works() {
    let agent = AgentBuilder::new(simple_text_model(1)).build();

    let history: Vec<Message> = Vec::new();

    let result = agent
        .prompt("test")
        .with_history(&history)
        .await
        .expect("prompt should succeed");

    assert_eq!(result, "hello from mock");

    // Note: The input history is not mutated. To get the updated history,
    // use `.extended_details()` and access `response.messages`.
}

/// Test 5: Multi-turn agent loop with tool calls — messages should contain the
/// full conversation: User → Assistant(tool_call) → User(tool_result) → Assistant(text).
#[tokio::test]
async fn multi_turn_messages_include_tool_calls() {
    let agent = AgentBuilder::new(tool_then_text_model())
        .tool(MockAddTool)
        .build();

    let resp = agent
        .prompt("What is 2 + 3?")
        .max_turns(5)
        .extended_details()
        .await
        .expect("prompt should succeed");

    assert_eq!(resp.output, "The answer is 5");

    let messages = resp.messages.expect("messages should be Some");

    // Expected sequence:
    // [0] User: "What is 2 + 3?"
    // [1] Assistant: ToolCall(add)
    // [2] User: ToolResult
    // [3] Assistant: "The answer is 5"
    assert_eq!(messages.len(), 4, "expected 4 messages, got: {messages:#?}");

    // [0] User prompt
    assert!(matches!(&messages[0], Message::User { .. }));

    // [1] Assistant with tool call
    match &messages[1] {
        Message::Assistant { content, .. } => {
            assert!(
                matches!(content.first(), AssistantContent::ToolCall(_)),
                "expected tool call, got: {content:?}"
            );
        }
        other => panic!("expected Assistant with tool call, got: {other:?}"),
    }

    // [2] User with tool result
    match &messages[2] {
        Message::User { content } => {
            assert!(
                matches!(content.first(), UserContent::ToolResult(_)),
                "expected tool result, got: {content:?}"
            );
        }
        other => panic!("expected User with tool result, got: {other:?}"),
    }

    // [3] Assistant with text
    match &messages[3] {
        Message::Assistant { content, .. } => match content.first() {
            AssistantContent::Text(t) => assert_eq!(t.text, "The answer is 5"),
            other => panic!("expected text, got: {other:?}"),
        },
        other => panic!("expected Assistant with text, got: {other:?}"),
    }

    // Usage should be aggregated across both turns
    assert_eq!(resp.usage.input_tokens, 35); // 15 + 20
    assert_eq!(resp.usage.output_tokens, 12); // 8 + 4
}

/// Test 6: `PromptResponse::new()` backward compatibility — 2-argument constructor
/// should still work, and `messages` should be `None`.
#[tokio::test]
async fn prompt_response_new_backward_compat() {
    use rig::agent::PromptResponse;

    let resp = PromptResponse::new("output text", Usage::new());

    assert_eq!(resp.output, "output text");
    assert!(resp.messages.is_none());
}

/// Test 6b: `PromptResponse` implements `Display`, delegating to `output`.
#[tokio::test]
async fn prompt_response_display_shows_output() {
    use rig::agent::PromptResponse;

    let resp = PromptResponse::new("the answer is 42", Usage::new());

    assert_eq!(format!("{resp}"), "the answer is 42");
    // Also works with format args
    assert_eq!(resp.to_string(), "the answer is 42");
}

/// Test 7: `PromptResponse::with_messages()` builder works correctly.
#[tokio::test]
async fn prompt_response_with_messages_builder() {
    use rig::agent::PromptResponse;

    let messages = vec![Message::user("hello"), Message::assistant("world")];

    let resp = PromptResponse::new("output", Usage::new()).with_messages(messages.clone());

    assert!(resp.messages.is_some());
    assert_eq!(resp.messages.as_ref().unwrap().len(), 2);
}

/// Test 8: MaxTurnsError still works — the error should contain the chat history.
/// This verifies the error path isn't broken by our changes.
#[tokio::test]
async fn max_turns_error_still_contains_history() {
    use rig::completion::PromptError;

    let agent = AgentBuilder::new(MockCompletionModel::new(
        (0..10).map(|_| always_tool_call_turn()),
    ))
    .tool(MockAddTool)
    .build();

    let result = agent
        .prompt("do something")
        .max_turns(2)
        .extended_details()
        .await;

    match result {
        Err(PromptError::MaxTurnsError {
            max_turns,
            chat_history,
            ..
        }) => {
            assert_eq!(max_turns, 2);
            // Chat history should have accumulated messages
            assert!(
                !chat_history.is_empty(),
                "chat_history in error should not be empty"
            );
        }
        Ok(_) => panic!("expected MaxTurnsError, got Ok"),
        Err(other) => panic!("expected MaxTurnsError, got: {other:?}"),
    }
}

/// Test 9: Extended details without `with_history()` — messages should still
/// be populated (this is the core feature: no need for &mut borrow).
#[tokio::test]
async fn extended_details_works_without_with_history() {
    let agent = AgentBuilder::new(tool_then_text_model())
        .tool(MockAddTool)
        .build();

    // Note: NO .with_history() call — this is the new use case
    let resp = agent
        .prompt("compute 2+3")
        .max_turns(5)
        .extended_details()
        .await
        .expect("prompt should succeed");

    let messages = resp
        .messages
        .expect("messages should be Some even without with_history()");

    // Should have full multi-turn history
    assert_eq!(messages.len(), 4);
    assert_eq!(resp.output, "The answer is 5");
}

/// Test 10: `Chat::chat` appends the prompt and response messages to the
/// caller-owned history.
#[tokio::test]
async fn chat_appends_prompt_and_assistant_to_history() {
    let agent = AgentBuilder::new(simple_text_model(2)).build();
    let mut history = Vec::<Message>::new();

    let output = agent
        .chat("hi", &mut history)
        .await
        .expect("chat should succeed");

    assert_eq!(output, "hello from mock");
    assert_eq!(
        history.len(),
        2,
        "expected chat to append [User, Assistant], got: {history:#?}"
    );

    match &history[0] {
        Message::User { content } => match content.first() {
            UserContent::Text(text) => assert_eq!(text.text, "hi"),
            other => panic!("expected text user content, got: {other:?}"),
        },
        other => panic!("expected User message, got: {other:?}"),
    }

    match &history[1] {
        Message::Assistant { content, .. } => match content.first() {
            AssistantContent::Text(text) => assert_eq!(text.text, "hello from mock"),
            other => panic!("expected text assistant content, got: {other:?}"),
        },
        other => panic!("expected Assistant message, got: {other:?}"),
    }

    let _ = agent
        .chat("again", &mut history)
        .await
        .expect("second chat should succeed");

    assert_eq!(
        history.len(),
        4,
        "expected history to grow to 4 messages, got: {history:#?}"
    );
}

/// Test 11: `Chat::chat` appends every message produced by a tool roundtrip.
#[tokio::test]
async fn chat_appends_tool_roundtrip_to_history() {
    let agent = AgentBuilder::new(tool_then_text_model())
        .tool(MockAddTool)
        .build();
    let mut history = Vec::<Message>::new();

    let output = agent
        .chat("What is 2 + 3?", &mut history)
        .await
        .expect("chat should succeed");

    assert_eq!(output, "The answer is 5");
    assert_eq!(
        history.len(),
        4,
        "expected chat to append [User, Assistant(tool), User(tool result), Assistant], got: {history:#?}"
    );
    assert!(matches!(&history[0], Message::User { .. }));

    match &history[1] {
        Message::Assistant { content, .. } => assert!(
            content
                .iter()
                .any(|content| matches!(content, AssistantContent::ToolCall(_))),
            "expected assistant tool call, got: {content:?}"
        ),
        other => panic!("expected Assistant with tool call, got: {other:?}"),
    }

    match &history[2] {
        Message::User { content } => assert!(
            content
                .iter()
                .any(|content| matches!(content, UserContent::ToolResult(_))),
            "expected user tool result, got: {content:?}"
        ),
        other => panic!("expected User with tool result, got: {other:?}"),
    }

    match &history[3] {
        Message::Assistant { content, .. } => match content.first() {
            AssistantContent::Text(text) => assert_eq!(text.text, "The answer is 5"),
            other => panic!("expected final assistant text, got: {other:?}"),
        },
        other => panic!("expected final Assistant, got: {other:?}"),
    }
}

/// Test 12: Multiple sequential prompts each return independent message histories.
#[tokio::test]
async fn sequential_prompts_have_independent_histories() {
    let agent = AgentBuilder::new(simple_text_model(2)).build();

    let resp1 = agent
        .prompt("first")
        .extended_details()
        .await
        .expect("first prompt should succeed");

    let resp2 = agent
        .prompt("second")
        .extended_details()
        .await
        .expect("second prompt should succeed");

    let msgs1 = resp1.messages.expect("messages should be Some");
    let msgs2 = resp2.messages.expect("messages should be Some");

    // Each should have exactly 2 messages (user + assistant)
    assert_eq!(msgs1.len(), 2);
    assert_eq!(msgs2.len(), 2);

    // First prompt's user message should be "first"
    match &msgs1[0] {
        Message::User { content } => match content.first() {
            UserContent::Text(t) => assert_eq!(t.text, "first"),
            other => panic!("unexpected: {other:?}"),
        },
        other => panic!("unexpected: {other:?}"),
    }

    // Second prompt's user message should be "second"
    match &msgs2[0] {
        Message::User { content } => match content.first() {
            UserContent::Text(t) => assert_eq!(t.text, "second"),
            other => panic!("unexpected: {other:?}"),
        },
        other => panic!("unexpected: {other:?}"),
    }
}
