use std::collections::HashSet;

use chrono::Local;
use uuid::Uuid;

use super::TaskStore;
use crate::ai::agent::task::{Task, TaskId};
use crate::ai::agent::{
    AIAgentExchange, AIAgentExchangeId, AIAgentOutput, AIAgentOutputMessage,
    AIAgentOutputMessageType, AIAgentOutputStatus, FinishedAIAgentOutput, MessageId, Shared,
    SubagentCall,
};
use crate::ai::llms::LLMId;
use crate::terminal::model::block::BlockId;

fn create_test_exchange() -> AIAgentExchange {
    AIAgentExchange {
        id: AIAgentExchangeId::new(),
        input: vec![],
        output_status: AIAgentOutputStatus::Streaming { output: None },
        added_message_ids: HashSet::new(),
        start_time: Local::now(),
        finish_time: None,
        time_to_first_token_ms: None,
        working_directory: None,
        model_id: LLMId::from(""),
        request_cost: None,
        coding_model_id: LLMId::from(""),
        cli_agent_model_id: LLMId::from(""),
        computer_use_model_id: LLMId::from(""),
        response_initiator: None,
    }
}

fn create_test_task_with_exchanges(exchange_count: usize) -> Task {
    let mut task = Task::new_optimistic_root();
    for _ in 0..exchange_count {
        task.append_exchange(create_test_exchange());
    }
    task
}

fn create_test_subtask_with_exchanges(exchange_count: usize) -> Task {
    let mut task = Task::new_optimistic_cli_agent_subtask(BlockId::new(), None);
    for _ in 0..exchange_count {
        task.append_exchange(create_test_exchange());
    }
    task
}

/// Creates an exchange with a finished output containing a subagent call to the given task_id.
fn create_exchange_with_subagent_call(subtask_id: &TaskId) -> AIAgentExchange {
    let output = AIAgentOutput {
        messages: vec![AIAgentOutputMessage {
            id: MessageId::new(Uuid::new_v4().to_string()),
            message: AIAgentOutputMessageType::Subagent(SubagentCall {
                task_id: subtask_id.to_string(),
                subagent_type: crate::ai::agent::SubagentType::Unknown,
            }),
            citations: vec![],
        }],
        citations: vec![],
        server_output_id: None,
        api_metadata_bytes: None,
        suggestions: None,
        telemetry_events: vec![],
        model_info: None,
        request_cost: None,
    };

    AIAgentExchange {
        id: AIAgentExchangeId::new(),
        input: vec![],
        output_status: AIAgentOutputStatus::Finished {
            finished_output: FinishedAIAgentOutput::Success {
                output: Shared::new(output),
            },
        },
        added_message_ids: HashSet::new(),
        start_time: Local::now(),
        finish_time: None,
        time_to_first_token_ms: None,
        working_directory: None,
        model_id: LLMId::from(""),
        request_cost: None,
        coding_model_id: LLMId::from(""),
        cli_agent_model_id: LLMId::from(""),
        computer_use_model_id: LLMId::from(""),
        response_initiator: None,
    }
}

#[test]
fn test_with_root_task() {
    let task = create_test_task_with_exchanges(3);
    let task_id = task.id().clone();
    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();

    let store = TaskStore::with_root_task(task);

    assert_eq!(store.task_count(), 1);
    assert_eq!(store.exchange_count(), 3);
    assert_eq!(store.root_task_id(), &task_id);
    assert_eq!(store.root_task().expect("task exists").id(), &task_id);
    assert_eq!(store.first_exchange().map(|e| e.id), Some(exchange_ids[0]));
    assert_eq!(store.latest_exchange().map(|e| e.id), Some(exchange_ids[2]));
}

#[test]
fn test_first_and_latest_exchange_o1() {
    let task = create_test_task_with_exchanges(5);
    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
    let store = TaskStore::with_root_task(task);

    // These should be O(1) operations
    let first = store.first_exchange().expect("has exchanges");
    let latest = store.latest_exchange().expect("has exchanges");

    assert_eq!(first.id, exchange_ids[0]);
    assert_eq!(latest.id, exchange_ids[4]);
}

#[test]
fn test_insert_subtask() {
    // Create root task with 1 exchange, then we'll add a subagent call exchange
    let root_task = create_test_task_with_exchanges(1);
    let root_task_id = root_task.id().clone();
    let mut store = TaskStore::with_root_task(root_task);

    // Create subtask with 1 exchange
    let subtask = create_test_subtask_with_exchanges(1);
    let subtask_id = subtask.id().clone();

    // Add exchange with subagent call to root task BEFORE inserting subtask
    let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
    store.append_exchange(&root_task_id, subagent_exchange);

    // Now insert the subtask - its exchanges should be included via the subagent call
    store.insert(subtask);

    assert_eq!(store.task_count(), 2);
    // 1 root exchange + 1 subagent call exchange + 1 subtask exchange = 3
    assert_eq!(store.exchange_count(), 3);
    assert!(store.get(&root_task_id).is_some());
    assert!(store.get(&subtask_id).is_some());
    assert!(store.contains(&subtask_id));
}

#[test]
fn test_unlinked_direct_cli_task_is_linearized_chronologically() {
    let base_time = Local::now();
    let mut root_task = Task::new_optimistic_root();
    let root_task_id = root_task.id().clone();

    let mut before_cli = create_test_exchange();
    before_cli.start_time = base_time;
    let before_cli_id = before_cli.id;
    root_task.append_exchange(before_cli);

    let mut after_cli = create_test_exchange();
    after_cli.start_time = base_time + chrono::Duration::seconds(2);
    let after_cli_id = after_cli.id;
    root_task.append_exchange(after_cli);

    let mut cli_task =
        Task::new_optimistic_cli_agent_subtask(BlockId::new(), Some(root_task_id.to_string()));
    let mut cli_exchange = create_test_exchange();
    cli_exchange.start_time = base_time + chrono::Duration::seconds(1);
    let cli_exchange_id = cli_exchange.id;
    cli_task.append_exchange(cli_exchange);

    let mut store = TaskStore::with_root_task(root_task);
    store.insert(cli_task);

    let exchange_ids = store
        .all_exchanges()
        .map(|exchange| exchange.id)
        .collect::<Vec<_>>();
    assert_eq!(
        exchange_ids,
        vec![before_cli_id, cli_exchange_id, after_cli_id]
    );
    assert_eq!(
        store.latest_exchange().map(|exchange| exchange.id),
        Some(after_cli_id)
    );
}

#[test]
fn test_remove_task() {
    let task = create_test_task_with_exchanges(3);
    let task_id = task.id().clone();
    let mut store = TaskStore::with_root_task(task);

    assert_eq!(store.task_count(), 1);
    assert_eq!(store.exchange_count(), 3);

    let removed = store.remove(&task_id);
    assert!(removed.is_some());
    assert_eq!(store.task_count(), 0);
    assert_eq!(store.exchange_count(), 0);
    assert!(store.get(&task_id).is_none());
}

#[test]
fn test_remove_nonexistent_task() {
    let task = create_test_task_with_exchanges(2);
    let mut store = TaskStore::with_root_task(task);

    let nonexistent_id = TaskId::new(Uuid::new_v4().to_string());
    let removed = store.remove(&nonexistent_id);
    assert!(removed.is_none());
    assert_eq!(store.task_count(), 1);
}

#[test]
fn test_all_exchanges_iteration() {
    let task = create_test_task_with_exchanges(4);
    let expected_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
    let store = TaskStore::with_root_task(task);

    let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();

    assert_eq!(actual_ids, expected_ids);
}

#[test]
fn test_all_exchanges_by_task() {
    let task = create_test_task_with_exchanges(3);
    let task_id = task.id().clone();
    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
    let store = TaskStore::with_root_task(task);

    let by_task = store.all_exchanges_by_task();
    assert_eq!(by_task.len(), 1);
    assert_eq!(by_task[0].0, task_id);
    assert_eq!(by_task[0].1.len(), 3);

    let actual_ids: Vec<_> = by_task[0].1.iter().map(|e| e.id).collect();
    assert_eq!(actual_ids, exchange_ids);
}

#[test]
fn test_set_root_task_replaces_old() {
    let task1 = create_test_task_with_exchanges(2);
    let task1_id = task1.id().clone();
    let mut store = TaskStore::with_root_task(task1);

    let task2 = create_test_task_with_exchanges(3);
    let task2_id = task2.id().clone();
    let task2_exchange_ids: Vec<_> = task2.exchanges().map(|e| e.id).collect();
    store.set_root_task(task2);

    assert_eq!(store.task_count(), 1);
    assert_eq!(store.exchange_count(), 3);
    // task1_id is now aliased to the new root task via optimistic_root_task_id
    assert!(store.get(&task1_id).is_some());
    assert!(store.get(&task2_id).is_some());
    assert_eq!(store.root_task_id(), &task2_id);

    let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(actual_ids, task2_exchange_ids);
}

#[test]
fn test_append_exchange() {
    let task = create_test_task_with_exchanges(2);
    let task_id = task.id().clone();
    let mut store = TaskStore::with_root_task(task);

    let new_exchange = create_test_exchange();
    let new_exchange_id = new_exchange.id;

    let result = store.append_exchange(&task_id, new_exchange);
    assert!(result);
    assert_eq!(store.exchange_count(), 3);

    // Verify the new exchange is accessible
    assert_eq!(store.latest_exchange().map(|e| e.id), Some(new_exchange_id));
}

#[test]
fn test_remove_task_exchange() {
    let task = create_test_task_with_exchanges(3);
    let task_id = task.id().clone();
    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
    let mut store = TaskStore::with_root_task(task);

    // First verify initial state
    assert_eq!(store.exchange_count(), 3);

    // Remove the middle exchange
    let removed = store.remove_task_exchange(&task_id, exchange_ids[1]);
    assert!(removed.is_some());
    assert_eq!(removed.unwrap().id, exchange_ids[1]);

    // After removal, we should have 2 exchanges
    assert_eq!(store.exchange_count(), 2);

    // Verify the remaining exchanges are correct
    let remaining_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(remaining_ids, vec![exchange_ids[0], exchange_ids[2]]);
}

#[test]
fn test_append_exchange_to_nonexistent_task() {
    let task = create_test_task_with_exchanges(1);
    let mut store = TaskStore::with_root_task(task);

    let nonexistent_id = TaskId::new(Uuid::new_v4().to_string());
    let result = store.append_exchange(&nonexistent_id, create_test_exchange());
    assert!(!result);
    assert_eq!(store.exchange_count(), 1);
}

#[test]
fn test_remove_exchange_from_nonexistent_task() {
    let task = create_test_task_with_exchanges(1);
    let exchange_id = task.exchanges().next().unwrap().id;
    let mut store = TaskStore::with_root_task(task);

    let nonexistent_task_id = TaskId::new(Uuid::new_v4().to_string());
    let result = store.remove_task_exchange(&nonexistent_task_id, exchange_id);
    assert!(result.is_none());
    assert_eq!(store.exchange_count(), 1);
}

#[test]
fn test_tasks_iteration() {
    let task1 = create_test_task_with_exchanges(2);
    let task1_id = task1.id().clone();
    let mut store = TaskStore::with_root_task(task1);

    let task2 = create_test_subtask_with_exchanges(1);
    let task2_id = task2.id().clone();
    store.insert(task2);

    let task_ids: HashSet<_> = store.tasks().map(|t| t.id().clone()).collect();
    assert_eq!(task_ids.len(), 2);
    assert!(task_ids.contains(&task1_id));
    assert!(task_ids.contains(&task2_id));
}

#[test]
fn test_multiple_tasks_exchange_order() {
    // Create root task with 1 exchange, then add subagent call exchange
    let root_task = create_test_task_with_exchanges(1);
    let root_task_id = root_task.id().clone();
    let first_root_exchange_id = root_task.exchanges().next().unwrap().id;
    let mut store = TaskStore::with_root_task(root_task);

    // Create subtask with 2 exchanges
    let subtask = create_test_subtask_with_exchanges(2);
    let subtask_id = subtask.id().clone();
    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();

    // Add subagent call exchange to root
    let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
    let subagent_exchange_id = subagent_exchange.id;
    store.append_exchange(&root_task_id, subagent_exchange);

    // Insert subtask - now linked via subagent call
    store.insert(subtask);

    // 1 root + 1 subagent call + 2 subtask = 4 exchanges
    assert_eq!(store.exchange_count(), 4);

    // First/last should still work
    assert_eq!(
        store.first_exchange().map(|e| e.id),
        Some(first_root_exchange_id)
    );
    assert_eq!(
        store.latest_exchange().map(|e| e.id),
        Some(subtask_exchange_ids[1])
    );

    // Order: root[0], root[subagent_call], subtask[0], subtask[1]
    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(all_ids.len(), 4);
    assert_eq!(all_ids[0], first_root_exchange_id);
    assert_eq!(all_ids[1], subagent_exchange_id);
    assert_eq!(all_ids[2], subtask_exchange_ids[0]);
    assert_eq!(all_ids[3], subtask_exchange_ids[1]);
}

#[test]
fn test_empty_task_handling() {
    let task = create_test_task_with_exchanges(0);
    let task_id = task.id().clone();
    let store = TaskStore::with_root_task(task);

    assert_eq!(store.task_count(), 1);
    assert_eq!(store.exchange_count(), 0);
    assert!(store.first_exchange().is_none());
    assert!(store.latest_exchange().is_none());
    assert!(store.get(&task_id).is_some());
}

#[test]
fn test_from_tasks() {
    use std::collections::HashMap;

    let task = create_test_task_with_exchanges(3);
    let task_id = task.id().clone();
    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();

    let mut tasks = HashMap::new();
    tasks.insert(task_id.clone(), task);

    let store = TaskStore::from_tasks(tasks, task_id.clone());

    assert_eq!(store.task_count(), 1);
    assert_eq!(store.exchange_count(), 3);
    assert_eq!(store.root_task_id(), &task_id);

    let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(actual_ids, exchange_ids);
}

#[test]
fn test_exchange_mut() {
    let task = create_test_task_with_exchanges(2);
    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
    let mut store = TaskStore::with_root_task(task);

    // Can find existing exchange
    let exchange = store.exchange_mut(exchange_ids[0]);
    assert!(exchange.is_some());
    assert_eq!(exchange.unwrap().id, exchange_ids[0]);

    // Returns None for non-existent exchange
    let fake_id = AIAgentExchangeId::new();
    assert!(store.exchange_mut(fake_id).is_none());
}

#[test]
fn test_modify_task_conditional_rebuild() {
    let task = create_test_task_with_exchanges(2);
    let task_id = task.id().clone();
    let original_exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
    let mut store = TaskStore::with_root_task(task);

    // Verify initial state
    assert_eq!(store.exchange_count(), 2);

    // modify_task with no exchange change should still work
    let result = store.modify_task(&task_id, |task| {
        assert_eq!(task.exchanges_len(), 2);
        "no change"
    });
    assert_eq!(result, Some("no change"));
    assert_eq!(store.exchange_count(), 2);

    // modify_task that adds an exchange should update the index
    let new_exchange = create_test_exchange();
    let new_exchange_id = new_exchange.id;
    store.modify_task(&task_id, |task| {
        task.append_exchange(new_exchange);
    });
    assert_eq!(store.exchange_count(), 3);
    assert_eq!(store.latest_exchange().map(|e| e.id), Some(new_exchange_id));

    // Verify all exchanges are in the index
    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(all_ids.len(), 3);
    assert_eq!(all_ids[0], original_exchange_ids[0]);
    assert_eq!(all_ids[1], original_exchange_ids[1]);
    assert_eq!(all_ids[2], new_exchange_id);
}

// =============================================================================
// Subtask Linearization Tests
// =============================================================================

#[test]
fn test_linearization_parent_with_one_subtask() {
    // Create a subtask first so we have its ID
    let subtask = create_test_subtask_with_exchanges(2);
    let subtask_id = subtask.id().clone();
    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();

    // Create root task: exchange1, exchange_with_subagent_call, exchange3
    let mut root_task = Task::new_optimistic_root();
    let exchange1 = create_test_exchange();
    let exchange1_id = exchange1.id;
    root_task.append_exchange(exchange1);

    let exchange_with_call = create_exchange_with_subagent_call(&subtask_id);
    let exchange_with_call_id = exchange_with_call.id;
    root_task.append_exchange(exchange_with_call);

    let exchange3 = create_test_exchange();
    let exchange3_id = exchange3.id;
    root_task.append_exchange(exchange3);

    // Build the store
    let mut store = TaskStore::with_root_task(root_task);
    store.insert(subtask);

    // Total exchanges: 3 root + 2 subtask = 5
    assert_eq!(store.exchange_count(), 5);

    // Expected order: root[0], root[1] (with subagent call), subtask[0], subtask[1], root[2]
    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(all_ids.len(), 5);
    assert_eq!(all_ids[0], exchange1_id);
    assert_eq!(all_ids[1], exchange_with_call_id);
    assert_eq!(all_ids[2], subtask_exchange_ids[0]);
    assert_eq!(all_ids[3], subtask_exchange_ids[1]);
    assert_eq!(all_ids[4], exchange3_id);

    // Verify first and last
    assert_eq!(store.first_exchange().map(|e| e.id), Some(exchange1_id));
    assert_eq!(store.latest_exchange().map(|e| e.id), Some(exchange3_id));
}

#[test]
fn test_linearization_nested_subtasks() {
    // Create nested subtask (grandchild) first
    let grandchild_subtask = create_test_subtask_with_exchanges(1);
    let grandchild_id = grandchild_subtask.id().clone();
    let grandchild_exchange_id = grandchild_subtask.exchanges().next().unwrap().id;

    // Create child subtask with a call to grandchild
    let mut child_subtask = Task::new_optimistic_cli_agent_subtask(BlockId::new(), None);
    let child_id = child_subtask.id().clone();

    let child_exchange1 = create_test_exchange();
    let child_exchange1_id = child_exchange1.id;
    child_subtask.append_exchange(child_exchange1);

    let child_call_to_grandchild = create_exchange_with_subagent_call(&grandchild_id);
    let child_call_exchange_id = child_call_to_grandchild.id;
    child_subtask.append_exchange(child_call_to_grandchild);

    // Create root task with a call to child
    let mut root_task = Task::new_optimistic_root();

    let root_exchange1 = create_test_exchange();
    let root_exchange1_id = root_exchange1.id;
    root_task.append_exchange(root_exchange1);

    let root_call_to_child = create_exchange_with_subagent_call(&child_id);
    let root_call_exchange_id = root_call_to_child.id;
    root_task.append_exchange(root_call_to_child);

    // Build the store
    let mut store = TaskStore::with_root_task(root_task);
    store.insert(child_subtask);
    store.insert(grandchild_subtask);

    // Total: 2 root + 2 child + 1 grandchild = 5
    assert_eq!(store.exchange_count(), 5);

    // Expected DFS order:
    // root[0], root[1] (calls child) -> child[0], child[1] (calls grandchild) -> grandchild[0]
    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(all_ids.len(), 5);
    assert_eq!(all_ids[0], root_exchange1_id);
    assert_eq!(all_ids[1], root_call_exchange_id);
    assert_eq!(all_ids[2], child_exchange1_id);
    assert_eq!(all_ids[3], child_call_exchange_id);
    assert_eq!(all_ids[4], grandchild_exchange_id);
}

#[test]
fn test_linearization_multiple_subtasks_same_parent() {
    // Create two subtasks
    let subtask1 = create_test_subtask_with_exchanges(1);
    let subtask1_id = subtask1.id().clone();
    let subtask1_exchange_id = subtask1.exchanges().next().unwrap().id;

    let subtask2 = create_test_subtask_with_exchanges(2);
    let subtask2_id = subtask2.id().clone();
    let subtask2_exchange_ids: Vec<_> = subtask2.exchanges().map(|e| e.id).collect();

    // Create root task with calls to both subtasks in separate exchanges
    let mut root_task = Task::new_optimistic_root();

    let call_to_subtask1 = create_exchange_with_subagent_call(&subtask1_id);
    let call_to_subtask1_id = call_to_subtask1.id;
    root_task.append_exchange(call_to_subtask1);

    let middle_exchange = create_test_exchange();
    let middle_exchange_id = middle_exchange.id;
    root_task.append_exchange(middle_exchange);

    let call_to_subtask2 = create_exchange_with_subagent_call(&subtask2_id);
    let call_to_subtask2_id = call_to_subtask2.id;
    root_task.append_exchange(call_to_subtask2);

    // Build the store
    let mut store = TaskStore::with_root_task(root_task);
    store.insert(subtask1);
    store.insert(subtask2);

    // Total: 3 root + 1 subtask1 + 2 subtask2 = 6
    assert_eq!(store.exchange_count(), 6);

    // Expected order:
    // root[0] (calls subtask1) -> subtask1[0], root[1], root[2] (calls subtask2) -> subtask2[0], subtask2[1]
    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
    assert_eq!(all_ids.len(), 6);
    assert_eq!(all_ids[0], call_to_subtask1_id);
    assert_eq!(all_ids[1], subtask1_exchange_id);
    assert_eq!(all_ids[2], middle_exchange_id);
    assert_eq!(all_ids[3], call_to_subtask2_id);
    assert_eq!(all_ids[4], subtask2_exchange_ids[0]);
    assert_eq!(all_ids[5], subtask2_exchange_ids[1]);
}

#[test]
fn test_all_exchanges_by_task_with_subtasks() {
    // Create a subtask
    let subtask = create_test_subtask_with_exchanges(2);
    let subtask_id = subtask.id().clone();
    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();

    // Create root task with a call to subtask
    let mut root_task = Task::new_optimistic_root();
    let root_id = root_task.id().clone();

    let root_exchange1 = create_test_exchange();
    let root_exchange1_id = root_exchange1.id;
    root_task.append_exchange(root_exchange1);

    let call_to_subtask = create_exchange_with_subagent_call(&subtask_id);
    let call_exchange_id = call_to_subtask.id;
    root_task.append_exchange(call_to_subtask);

    let root_exchange3 = create_test_exchange();
    let root_exchange3_id = root_exchange3.id;
    root_task.append_exchange(root_exchange3);

    // Build the store
    let mut store = TaskStore::with_root_task(root_task);
    store.insert(subtask);

    // Check all_exchanges_by_task grouping
    let by_task = store.all_exchanges_by_task();

    // Should have 3 groups: root[0-1], subtask[0-1], root[2]
    assert_eq!(by_task.len(), 3);

    // First group: root task's first two exchanges
    assert_eq!(by_task[0].0, root_id);
    assert_eq!(by_task[0].1.len(), 2);
    assert_eq!(by_task[0].1[0].id, root_exchange1_id);
    assert_eq!(by_task[0].1[1].id, call_exchange_id);

    // Second group: subtask's exchanges
    assert_eq!(by_task[1].0, subtask_id);
    assert_eq!(by_task[1].1.len(), 2);
    assert_eq!(by_task[1].1[0].id, subtask_exchange_ids[0]);
    assert_eq!(by_task[1].1[1].id, subtask_exchange_ids[1]);

    // Third group: root task's last exchange
    assert_eq!(by_task[2].0, root_id);
    assert_eq!(by_task[2].1.len(), 1);
    assert_eq!(by_task[2].1[0].id, root_exchange3_id);
}

#[test]
fn test_exchange_by_id_resolves_subtask_exchange() {
    // The index spans all tasks, not just the root, so exchanges that live in a
    // subtask must be resolvable by id.
    let root_task = create_test_task_with_exchanges(1);
    let root_task_id = root_task.id().clone();
    let mut store = TaskStore::with_root_task(root_task);

    let subtask = create_test_subtask_with_exchanges(2);
    let subtask_id = subtask.id().clone();
    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();

    let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
    store.append_exchange(&root_task_id, subagent_exchange);
    store.insert(subtask);

    for id in &subtask_exchange_ids {
        assert_eq!(store.exchange_by_id(*id).map(|e| e.id), Some(*id));
    }
}

#[test]
fn test_exchange_by_id_after_remove_task() {
    let root_task = create_test_task_with_exchanges(1);
    let root_task_id = root_task.id().clone();
    let root_exchange_id = root_task.exchanges().next().unwrap().id;
    let mut store = TaskStore::with_root_task(root_task);

    let subtask = create_test_subtask_with_exchanges(2);
    let subtask_id = subtask.id().clone();
    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();

    let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
    store.append_exchange(&root_task_id, subagent_exchange);
    store.insert(subtask);

    // Sanity: the subtask's exchanges resolve before removal.
    assert!(store.exchange_by_id(subtask_exchange_ids[0]).is_some());

    store.remove(&subtask_id);

    // The removed task's exchanges are no longer resolvable.
    for id in &subtask_exchange_ids {
        assert!(store.exchange_by_id(*id).is_none());
    }
    // The surviving root exchange still resolves.
    assert_eq!(
        store.exchange_by_id(root_exchange_id).map(|e| e.id),
        Some(root_exchange_id)
    );
}

#[test]
fn test_exchange_by_id_after_remove_task_exchange_index_shift() {
    let task = create_test_task_with_exchanges(3);
    let task_id = task.id().clone();
    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
    let mut store = TaskStore::with_root_task(task);

    // Remove the middle exchange, which shifts the index of everything after it.
    store.remove_task_exchange(&task_id, exchange_ids[1]);

    // The removed id no longer resolves.
    assert!(store.exchange_by_id(exchange_ids[1]).is_none());

    // The exchange that followed it (now at a different index) still resolves to itself.
    assert_eq!(
        store.exchange_by_id(exchange_ids[2]).map(|e| e.id),
        Some(exchange_ids[2])
    );
    // The exchange before it is unaffected.
    assert_eq!(
        store.exchange_by_id(exchange_ids[0]).map(|e| e.id),
        Some(exchange_ids[0])
    );
}

#[test]
fn test_exchange_by_id_after_modify_task_append() {
    let task = create_test_task_with_exchanges(2);
    let task_id = task.id().clone();
    let mut store = TaskStore::with_root_task(task);

    let new_exchange = create_test_exchange();
    let new_exchange_id = new_exchange.id;
    store.modify_task(&task_id, |task| {
        task.append_exchange(new_exchange);
    });

    // The newly appended exchange is found via the rebuilt index.
    assert_eq!(
        store.exchange_by_id(new_exchange_id).map(|e| e.id),
        Some(new_exchange_id)
    );
}

#[test]
fn test_exchange_by_id_after_set_root_task() {
    let task1 = create_test_task_with_exchanges(2);
    let task1_exchange_ids: Vec<_> = task1.exchanges().map(|e| e.id).collect();
    let mut store = TaskStore::with_root_task(task1);

    let task2 = create_test_task_with_exchanges(3);
    let task2_exchange_ids: Vec<_> = task2.exchanges().map(|e| e.id).collect();
    store.set_root_task(task2);

    // The old root's exchanges no longer resolve.
    for id in &task1_exchange_ids {
        assert!(store.exchange_by_id(*id).is_none());
    }
    // The new root's exchanges resolve.
    for id in &task2_exchange_ids {
        assert_eq!(store.exchange_by_id(*id).map(|e| e.id), Some(*id));
    }
}