use rig::completion::message::{ToolResultContent, UserContent};
use rig::completion::{AssistantContent, Message};

const MAX_TOOL_RESULT_CHARS: usize = 2000;

/// Best-effort reconstruction of research material from a chat history left
/// behind when the researcher hit its turn budget mid-investigation: whatever
/// prose the model wrote between tool calls, plus the raw text of every tool
/// result (search snippets, fetched pages), truncated so one huge page can't
/// crowd out everything else that was found.
pub(crate) fn partial_findings_from_history(chat_history: &[Message]) -> String {
    let sections: Vec<String> = chat_history
        .iter()
        .flat_map(|message| {
            extract_assistant_text(message)
                .into_iter()
                .chain(tool_result_text(message, MAX_TOOL_RESULT_CHARS))
        })
        .collect();

    match sections.is_empty() {
        true => "The researcher exhausted its turn budget before gathering any usable evidence."
            .to_string(),
        false => sections.join("\n"),
    }
}

/// Plain-text blocks from an assistant message, if any.
fn extract_assistant_text(message: &Message) -> Vec<String> {
    let Message::Assistant { content, .. } = message else {
        return Vec::new();
    };
    content
        .iter()
        .filter_map(|item: &AssistantContent| match item {
            AssistantContent::Text(text) => Some(text.text().to_string()),
            _ => None,
        })
        .collect()
}

/// Plain-text tool results attached to a user message (that's where rig
/// places them), each truncated to `max_chars` so one huge fetched page
/// can't crowd out everything else that was found.
fn tool_result_text(message: &Message, max_chars: usize) -> Vec<String> {
    let Message::User { content } = message else {
        return Vec::new();
    };

    content
        .iter()
        .filter_map(|item: &UserContent| match item {
            UserContent::ToolResult(tool_result) => Some(tool_result),
            _ => None,
        })
        .flat_map(|tool_result| tool_result.content.iter())
        .filter_map(
            |tool_result_content: &ToolResultContent| match tool_result_content {
                ToolResultContent::Text(text) => Some(truncate(text.text(), max_chars)),
                _ => None,
            },
        )
        .collect()
}

fn truncate(text: &str, max_chars: usize) -> String {
    let mut result = text.to_string();
    if result.len() > max_chars {
        result.truncate(max_chars);
        result.push_str(" ...[truncated]");
    }
    result
}

/// Chronological transcript of a partial research run, annotated with tool
/// calls (so a fetch's URL or a search's query stays attached to its
/// result) rather than just the bare result text `partial_findings_from_history`
/// collects — the summarizer agent needs that context to attribute facts to
/// the right source.
pub(crate) fn annotated_transcript_from_history(chat_history: &[Message]) -> String {
    chat_history
        .iter()
        .flat_map(transcript_lines)
        .collect::<Vec<String>>()
        .join("\n")
}

fn transcript_lines(message: &Message) -> Vec<String> {
    match message {
        Message::Assistant { content, .. } => content
            .iter()
            .filter_map(|item| match item {
                AssistantContent::Text(text) => Some(format!("Note: {}", text.text())),
                AssistantContent::ToolCall(call) => Some(format!(
                    "Called {}({})",
                    call.function.name, call.function.arguments
                )),
                _ => None,
            })
            .collect(),
        Message::User { content } => content
            .iter()
            .filter_map(|item| match item {
                UserContent::ToolResult(result) => Some(result),
                _ => None,
            })
            .flat_map(|result| result.content.iter())
            .filter_map(|part| match part {
                ToolResultContent::Text(text) => Some(format!(
                    "Result: {}",
                    truncate(text.text(), MAX_TOOL_RESULT_CHARS)
                )),
                _ => None,
            })
            .collect(),
        Message::System { .. } => Vec::new(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn truncate_leaves_short_text_untouched() {
        assert_eq!(truncate("hello", 10), "hello");
    }

    #[test]
    fn truncate_leaves_exact_length_text_untouched() {
        assert_eq!(truncate("hello", 5), "hello");
    }

    #[test]
    fn truncate_cuts_long_text_and_marks_it() {
        assert_eq!(truncate("hello world", 5), "hello ...[truncated]");
    }

    #[test]
    fn assistant_text_extracts_text_blocks() {
        let message = Message::assistant("found it");
        assert_eq!(
            extract_assistant_text(&message),
            vec!["found it".to_string()]
        );
    }

    #[test]
    fn assistant_text_ignores_tool_calls() {
        let message = Message::Assistant {
            id: None,
            content: vec![AssistantContent::tool_call(
                "call-1",
                "search_web",
                serde_json::json!({ "query": "test"}),
            )],
        };
        assert!(extract_assistant_text(&message).is_empty());
    }

    #[test]
    fn assistant_text_ignores_non_assistant_messages() {
        assert!(extract_assistant_text(&Message::user("hi")).is_empty());
        assert!(extract_assistant_text(&Message::system("be careful")).is_empty());
    }

    #[test]
    fn tool_result_text_extracts_and_truncates() {
        let short = Message::tool_result("call-1", "call-1-name", "short result");
        assert_eq!(
            tool_result_text(&short, 100),
            vec!["short result".to_string()]
        );

        let long = Message::tool_result("call-2", "call-2-result", "0123456789");
        assert_eq!(
            tool_result_text(&long, 5),
            vec!["01234 ...[truncated]".to_string()]
        );
    }

    #[test]
    fn tool_result_text_ignores_non_tool_result_content() {
        assert!(tool_result_text(&Message::user("plain text, no tool result"), 100).is_empty());
        assert!(tool_result_text(&Message::assistant("also ignored"), 100).is_empty());
    }

    #[test]
    fn partial_findings_from_history_falls_back_when_nothing_usable() {
        let history = vec![
            Message::system("preamble"),
            Message::Assistant {
                id: None,
                content: vec![AssistantContent::tool_call(
                    "call-1",
                    "search_web",
                    serde_json::json!({ "query": "test" }),
                )],
            },
        ];

        assert_eq!(
            partial_findings_from_history(&history),
            "The researcher exhausted its turn budget before gathering any usable evidence."
        );
    }

    #[test]
    fn partial_findings_from_history_collects_assistant_text_and_tool_results_in_order() {
        let history = vec![
            Message::assistant("Checking sources..."),
            Message::tool_result("call-1", "call-1-name", "Result A [1]"),
            Message::assistant("Cross-checking..."),
            Message::tool_result("call-2", "call-2-name", "Result B [2]"),
        ];

        assert_eq!(
            partial_findings_from_history(&history),
            "Checking sources...\nResult A [1]\nCross-checking...\nResult B [2]"
        );
    }

    #[test]
    fn transcript_lines_prefixes_assistant_notes() {
        let message = Message::assistant("still checking this claim");
        assert_eq!(
            transcript_lines(&message),
            vec!["Note: still checking this claim".to_string()]
        );
    }

    #[test]
    fn transcript_lines_renders_tool_calls_with_their_arguments() {
        let message = Message::Assistant {
            id: None,
            content: vec![AssistantContent::tool_call(
                "call-1",
                "search_web",
                serde_json::json!({ "query": "uruguay senior engineers" }),
            )],
        };
        assert_eq!(
            transcript_lines(&message),
            vec!["Called search_web({\"query\":\"uruguay senior engineers\"})".to_string()]
        );
    }

    #[test]
    fn transcript_lines_prefixes_and_truncates_tool_results() {
        let short = Message::tool_result("call-1", "call-1-name", "found via search [1]");
        assert_eq!(
            transcript_lines(&short),
            vec!["Result: found via search [1]".to_string()]
        );
    }

    #[test]
    fn transcript_lines_ignores_system_messages() {
        assert!(transcript_lines(&Message::system("be thorough")).is_empty());
    }

    #[test]
    fn annotated_transcript_from_history_joins_calls_results_and_notes_in_order() {
        let history = vec![
            Message::Assistant {
                id: None,
                content: vec![AssistantContent::tool_call(
                    "call-1",
                    "search_web",
                    serde_json::json!({ "query": "test" }),
                )],
            },
            Message::tool_result("call-1", "call-1-name", "1. Example\n   https://example.com\n   snippet"),
            Message::assistant("that source looks solid"),
        ];

        assert_eq!(
            annotated_transcript_from_history(&history),
            "Called search_web({\"query\":\"test\"})\n\
             Result: 1. Example\n   https://example.com\n   snippet\n\
             Note: that source looks solid"
        );
    }

    #[test]
    fn annotated_transcript_from_history_is_empty_with_no_usable_content() {
        assert_eq!(
            annotated_transcript_from_history(&[Message::system("preamble")]),
            ""
        );
    }
}