//! Model-facing output extraction.
//!
//! Tool outputs carry both structured data (for agent/client logic) and
//! a model-facing representation as MCP content blocks. [`ToolOutput`]
//! is the trait that the runtime uses to extract the model-facing part.
//!
//! The default [`ToolOutput`] implementation serialises the output
//! to JSON, then walks the structure looking for embedded
//! [`ContentBlock`]-shaped values (images, resources). These are
//! promoted to proper content block types; everything else becomes
//! [`ContentBlock::Text`].
//!
//! Use [`extract_content_blocks`] directly when you need the same
//! conversion on an arbitrary [`serde_json::Value`].
//!
//! # Example — custom model output
//!
//! ```rust
//! use serde::Serialize;
//! use xai_tool_runtime::render::ToolOutput;
//! use xai_tool_runtime::ContentBlock;
//!
//! #[derive(Serialize)]
//! struct BashOutput {
//!     stdout: String,
//!     exit_code: i32,
//!     model_output: Vec<ContentBlock>,
//! }
//!
//! impl ToolOutput for BashOutput {
//!     fn model_output(&self) -> Vec<ContentBlock> {
//!         self.model_output.clone()
//!     }
//! }
//! ```
//!
//! # Example — default (automatic MCP extraction)
//!
//! Types that don't override `model_output()` get automatic extraction.
//! Embedded images and resources are promoted; the rest is JSON text:
//!
//! ```rust
//! use serde::Serialize;
//! use xai_tool_runtime::render::ToolOutput;
//!
//! #[derive(Serialize)]
//! struct SimpleOutput { answer: String }
//! impl ToolOutput for SimpleOutput {}
//! // model sees: ContentBlock::Text { text: r#"{"answer":"..."}"# }
//! ```

use std::sync::Arc;

use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::tool::ContentBlock;

/// Unified trait for typed tool outputs.
///
/// Combines model-facing content extraction and optional
/// chat-completion response generation into a single trait.
pub trait ToolOutput: Serialize {
    /// Returns the model-facing content blocks for this output.
    ///
    /// Return an empty `Vec` to signal "use automatic extraction" — the
    /// runtime will call [`extract_content_blocks`] on the serialised
    /// JSON value instead.
    fn model_output(&self) -> Vec<ContentBlock> {
        Vec::new()
    }

    /// Build a chat-completion response frame from this tool output (sent to client),
    /// if applicable.  Returns `None` by default.
    fn chat_completion_output(&self) -> Option<ToolChatCompletionResponse> {
        None
    }
}

/// Blanket impl so `serde_json::Value` can be used directly as a
/// `Tool::Output` (handy for stub/test tools and pass-through proxies).
impl ToolOutput for Value {}

/// `String` is a common output type for simple tools.
impl ToolOutput for String {}

/// Lets `xai_tool_types::TaskOutputOutput` be used directly as a `Tool::Output`
/// (handy for stub/test tools and pass-through proxies).
impl ToolOutput for xai_tool_types::TaskOutputOutput {}

/// Lets `xai_tool_types::SubagentCompletedOutput` be used directly as a
/// `Tool::Output` (the `task` tool's structured completion output).
impl ToolOutput for xai_tool_types::SubagentCompletedOutput {}

/// Lets `xai_tool_types::KillTaskOutput` be used directly as a `Tool::Output`
/// (the `kill_task` tool's typed result / not-found output).
impl ToolOutput for xai_tool_types::KillTaskOutput {}

/// Delegate through `Box<T>` so boxed outputs (e.g. large response
/// structs) work without a manual impl.
impl<T: ToolOutput + Serialize + ?Sized> ToolOutput for Box<T> {
    fn model_output(&self) -> Vec<ContentBlock> {
        (**self).model_output()
    }

    fn chat_completion_output(&self) -> Option<ToolChatCompletionResponse> {
        (**self).chat_completion_output()
    }
}

/// Minimal representation of a chat-completion response streamed to the
/// frontend.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ToolChatCompletionResponse {
    /// The main completion payload.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<ToolChatCompletion>,
    /// Structured stream error (e.g. rate-limit, tool failure).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream_error: Option<ToolStreamError>,
}

/// Minimal chat completion response for tool result (sent to client).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ToolChatCompletion {
    /// Always `"assistant"`.
    #[serde(default)]
    pub sender: String,
    /// Text body of the response.
    #[serde(default)]
    pub message: String,
    /// Tag discriminator: `"final"`, `"raw_function_result"`,
    /// `"tool_usage_card"`, `"tool_partial_output"`, etc.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message_tag: Option<String>,
    /// Identifies the tool-usage card this result belongs to.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_usage_card_id: Option<String>,
    /// JSON-encoded card attachment (images, render cards, files).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub card_attachment: Option<String>,
    /// Media generation type: `"image_gen"`, `"video_gen"`, etc.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub media_gen_type: Option<String>,
    /// Code execution result.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub code_execution_result: Option<ToolCodeExecutionResult>,
    /// Catch-all for additional fields the tool wants to set. Merged
    /// into the proto `ChatCompletion` by the downstream converter.
    #[serde(flatten)]
    pub extra: serde_json::Map<String, Value>,
}

/// Lightweight code-execution result carried on the completion.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ToolCodeExecutionResult {
    #[serde(default)]
    pub stdout: String,
    #[serde(default)]
    pub stderr: String,
    #[serde(default)]
    pub exit_code: i32,
    #[serde(default)]
    pub command_timed_out: bool,
}

/// Structured stream error returned alongside the completion.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ToolStreamError {
    pub message: String,
    /// Opaque typed-error payload. The downstream chat layer
    /// deserialises this into the concrete proto enum variant.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub typed_error: Option<Value>,
}

/// Extract MCP-compatible [`ContentBlock`]s from a serialised JSON value.
///
/// Strategies are tried in order — first match wins:
///
/// | # | Shape | Result |
/// |---|-------|--------|
/// | 1 | Value is itself a `ContentBlock` (`{"type":"text",…}`) | `vec![block]` |
/// | 2 | Array containing ≥ 1 `ContentBlock` | each element: block or text |
/// | 3 | Object with `"content": [...]` (MCP `CallToolResult`) | `structuredContent` (if any) as JSON text, followed by the content array |
/// | 4 | Object with mixed fields | block-shaped fields extracted, rest as JSON text |
/// | 5 | Anything else | `ContentBlock::Text` with the stringified value |
pub fn extract_content_blocks(value: &Value) -> Vec<ContentBlock> {
    // 1. Value IS a single ContentBlock.
    if let Some(block) = try_parse_block(value) {
        return vec![block];
    }

    // 2. Array: convert each element (block-shaped -> block, else -> text).
    //    Only enter this path when at least one element looks like a
    //    ContentBlock so plain arrays like [1,2,3] fall through to text.
    if let Some(arr) = value.as_array()
        && !arr.is_empty()
        && arr.iter().any(looks_like_content_block)
    {
        return arr.iter().map(value_to_block).collect();
    }

    if let Some(obj) = value.as_object() {
        // grok-build `ToolRunResult`: the model sees `prompt_text` (reminders
        // appended), never a JSON dump of the structured result.
        if let Some(Value::String(prompt_text)) = obj.get("prompt_text")
            && obj.contains_key("output")
            && obj.contains_key("effective_tool_name")
        {
            return vec![ContentBlock::Text {
                text: prompt_text.clone(),
            }];
        }

        // 3. Object with a `"content"` array -> the standard MCP
        //    CallToolResult shape.
        if let Some(Value::Array(arr)) = obj.get("content")
            && !arr.is_empty()
            && arr.iter().any(looks_like_content_block)
        {
            // Surface `structuredContent` so IDs/handles the server
            // expects the model to round-trip aren't dropped.
            let structured = obj
                .get("structuredContent")
                .filter(|v| !v.is_null())
                .map(|v| ContentBlock::Text {
                    text: v.to_string(),
                });
            let mut blocks = Vec::with_capacity(arr.len() + structured.is_some() as usize);
            blocks.extend(structured);
            blocks.extend(arr.iter().map(value_to_block));
            return blocks;
        }

        // 4. Mixed object -> pull block-shaped field values out; collect
        //    the remaining fields into a single JSON text block.
        let mut extracted = Vec::new();
        let mut remainder = serde_json::Map::new();

        for (key, val) in obj {
            match classify_field(val) {
                FieldShape::Block(block) => extracted.push(block),
                FieldShape::Blocks(blocks) => extracted.extend(blocks),
                FieldShape::Other => {
                    remainder.insert(key.clone(), val.clone());
                }
            }
        }

        if !extracted.is_empty() {
            let mut result = Vec::new();
            if !remainder.is_empty() {
                result.push(ContentBlock::Text {
                    text: Value::Object(remainder).to_string(),
                });
            }
            result.extend(extracted);
            return result;
        }
    }

    // 5. Fallback -> render the whole value as text.
    vec![value_to_block(value)]
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/// The `ContentBlock` enum is `#[serde(tag = "type", rename_all =
/// "snake_case")]`, so a JSON object can only be a content block when
/// it has `"type"` set to one of these three values.
const CONTENT_BLOCK_TYPES: &[&str] = &["text", "image", "resource"];

/// Cheap check: could `value` plausibly deserialise as a
/// [`ContentBlock`]?  Only objects with a `"type"` field whose value
/// is one of the known discriminators pass.  This avoids a full
/// `from_value(clone())` on the vast majority of values.
fn looks_like_content_block(value: &Value) -> bool {
    value
        .as_object()
        .and_then(|obj| obj.get("type"))
        .and_then(Value::as_str)
        .is_some_and(|t| CONTENT_BLOCK_TYPES.contains(&t))
}

/// Try to parse `value` as a [`ContentBlock`].  Returns `None`
/// immediately when the value doesn't pass the cheap
/// [`looks_like_content_block`] check, avoiding clone + full
/// deserialisation for non-matching shapes.
fn try_parse_block(value: &Value) -> Option<ContentBlock> {
    if !looks_like_content_block(value) {
        return None;
    }
    serde_json::from_value::<ContentBlock>(value.clone()).ok()
}

/// Result of inspecting a single object field value.
enum FieldShape {
    /// The field value IS a single `ContentBlock`.
    Block(ContentBlock),
    /// The field value is an array where *every* element is a `ContentBlock`.
    Blocks(Vec<ContentBlock>),
    /// The field value does not look like block content.
    Other,
}

/// Classify a field value as block content.
///
/// Conservative for arrays: all elements must deserialise as
/// `ContentBlock`; mixed arrays go to `Other` so ambiguous data
/// (e.g. `"scores": [0.9, 0.8]`) is not silently dropped.
fn classify_field(value: &Value) -> FieldShape {
    // Single block.
    if let Some(block) = try_parse_block(value) {
        return FieldShape::Block(block);
    }
    // Array of blocks — strict: every element must parse.
    if let Some(arr) = value.as_array()
        && !arr.is_empty()
        && arr.iter().all(looks_like_content_block)
    {
        let blocks: Result<Vec<ContentBlock>, _> = arr
            .iter()
            .map(|v| serde_json::from_value::<ContentBlock>(v.clone()))
            .collect();
        if let Ok(blocks) = blocks {
            return FieldShape::Blocks(blocks);
        }
    }
    FieldShape::Other
}

/// Convert a single `Value` to a `ContentBlock`.
///
/// Uses the cheap [`try_parse_block`] check first; on failure wraps
/// the value as `ContentBlock::Text`.  Strings are used verbatim (no
/// extra JSON quoting); all other types go through `Value::to_string`.
fn value_to_block(value: &Value) -> ContentBlock {
    try_parse_block(value).unwrap_or_else(|| ContentBlock::Text {
        text: match value {
            Value::String(s) => s.clone(),
            other => other.to_string(),
        },
    })
}

// ---------------------------------------------------------------------------
// Type-erased extractor (used by the toolbox registry)
// ---------------------------------------------------------------------------

/// Type-erased model output extractor.
pub type ModelOutputExtractor = Arc<dyn Fn(&Value) -> Option<Vec<ContentBlock>> + Send + Sync>;

/// Build a [`ModelOutputExtractor`] for a concrete output type.
pub fn extractor_for<T>() -> ModelOutputExtractor
where
    T: ToolOutput + serde::de::DeserializeOwned + 'static,
{
    Arc::new(|value: &Value| {
        serde_json::from_value::<T>(value.clone())
            .ok()
            .map(|output| {
                let blocks = output.model_output();
                if blocks.is_empty() {
                    extract_content_blocks(value)
                } else {
                    blocks
                }
            })
    })
}

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

    // ── ToolOutput with custom override ─────────────────────────────

    #[derive(Serialize)]
    struct FakeOutput {
        blocks: Vec<ContentBlock>,
    }

    impl ToolOutput for FakeOutput {
        fn model_output(&self) -> Vec<ContentBlock> {
            self.blocks.clone()
        }
    }

    #[test]
    fn custom_text_block() {
        let o = FakeOutput {
            blocks: vec![ContentBlock::Text {
                text: "hello".into(),
            }],
        };
        assert_eq!(o.model_output().len(), 1);
        assert_eq!(
            o.model_output()[0],
            ContentBlock::Text {
                text: "hello".into()
            }
        );
    }

    #[test]
    fn custom_multimodal() {
        let o = FakeOutput {
            blocks: vec![
                ContentBlock::Text {
                    text: "result:".into(),
                },
                ContentBlock::Image {
                    mime_type: "image/png".into(),
                    data: "iVBOR...".into(),
                    media_id: None,
                    filename: None,
                    path: None,
                    metadata: Default::default(),
                },
            ],
        };
        assert_eq!(o.model_output().len(), 2);
    }

    // ── ToolOutput default → empty (runtime fills via extract) ──────

    #[test]
    fn default_model_output_returns_empty() {
        #[derive(Serialize)]
        struct Plain {
            value: u32,
        }
        impl ToolOutput for Plain {}

        // Default signals "use automatic extraction" by returning empty.
        assert!(Plain { value: 42 }.model_output().is_empty());
    }

    #[test]
    fn runtime_fills_empty_model_output_via_extract() {
        // Simulates what the ToolDyn blanket does: serialise once,
        // then extract_content_blocks on the Value.
        #[derive(Serialize)]
        struct Plain {
            value: u32,
        }
        impl ToolOutput for Plain {}

        let p = Plain { value: 42 };
        let value = serde_json::to_value(&p).unwrap();
        let custom = p.model_output();
        let blocks = if custom.is_empty() {
            extract_content_blocks(&value)
        } else {
            custom
        };
        assert_eq!(blocks.len(), 1);
        assert_eq!(
            blocks[0],
            ContentBlock::Text {
                text: r#"{"value":42}"#.into(),
            }
        );
    }

    // ── extract_content_blocks unit tests ──────────────────────────

    // Strategy 1: single ContentBlock
    #[test]
    fn extract_single_text_block() {
        let v = json!({"type": "text", "text": "hi"});
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks, vec![ContentBlock::Text { text: "hi".into() }]);
    }

    #[test]
    fn extract_single_image_block() {
        let v = json!({"type": "image", "mime_type": "image/png", "data": "abc"});
        let blocks = extract_content_blocks(&v);
        assert_eq!(
            blocks,
            vec![ContentBlock::Image {
                mime_type: "image/png".into(),
                data: "abc".into(),
                media_id: None,
                filename: None,
                path: None,
                metadata: Default::default(),
            }]
        );
    }

    #[test]
    fn extract_single_resource_block() {
        let v = json!({"type": "resource", "uri": "file:///x"});
        let blocks = extract_content_blocks(&v);
        assert_eq!(
            blocks,
            vec![ContentBlock::Resource {
                uri: "file:///x".into(),
                mime_type: None,
                text: None,
            }]
        );
    }

    #[test]
    fn extract_mcp_image_with_camel_case() {
        let v = json!({"type": "image", "mimeType": "image/png", "data": "abc"});
        let blocks = extract_content_blocks(&v);
        assert_eq!(
            blocks,
            vec![ContentBlock::Image {
                mime_type: "image/png".into(),
                data: "abc".into(),
                media_id: None,
                filename: None,
                path: None,
                metadata: Default::default(),
            }]
        );
    }

    // Strategy 2: array of blocks
    #[test]
    fn extract_array_of_blocks() {
        let v = json!([
            {"type": "text", "text": "a"},
            {"type": "image", "mime_type": "image/png", "data": "b"},
        ]);
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 2);
        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "a"));
        assert!(matches!(&blocks[1], ContentBlock::Image { .. }));
    }

    #[test]
    fn extract_mixed_array_promotes_non_blocks_to_text() {
        let v = json!([
            {"type": "text", "text": "a"},
            42,
            "raw string",
        ]);
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 3);
        assert_eq!(blocks[0], ContentBlock::Text { text: "a".into() });
        assert_eq!(blocks[1], ContentBlock::Text { text: "42".into() });
        assert_eq!(
            blocks[2],
            ContentBlock::Text {
                text: "raw string".into()
            }
        );
    }

    #[test]
    fn plain_array_falls_through_to_text() {
        // No block-shaped elements → single text fallback.
        let v = json!([1, 2, 3]);
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 1);
        assert_eq!(
            blocks[0],
            ContentBlock::Text {
                text: "[1,2,3]".into()
            }
        );
    }

    // Strategy 3: object with "content" key
    #[test]
    fn extract_content_field() {
        let v = json!({
            "is_error": false,
            "content": [
                {"type": "text", "text": "summary"},
                {"type": "image", "mime_type": "image/png", "data": "b64"},
            ],
        });
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 2);
        assert_eq!(
            blocks[0],
            ContentBlock::Text {
                text: "summary".into()
            }
        );
        assert!(matches!(blocks[1], ContentBlock::Image { .. }));
    }

    // Strategy 3 + structuredContent (MCP CallToolResult)
    #[test]
    fn extract_content_with_structured_content_surfaces_id() {
        let v = json!({
            "content": [
                {"type": "resource", "uri": "ui://tldraw/canvas",
                 "mimeType": "text/html", "text": "<html>…</html>"},
            ],
            "structuredContent": {"drawing_id": "abc123", "title": "sketch"},
            "isError": false,
            "_meta": {"ui": {"resourceUri": "ui://tldraw/canvas"}},
        });
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 2, "expected structured + resource");
        // structuredContent rendered as JSON text, ahead of content.
        let ContentBlock::Text { text } = &blocks[0] else {
            panic!("expected first block to be Text, got {:?}", blocks[0]);
        };
        assert!(
            text.contains("abc123"),
            "structuredContent must surface drawing_id, got: {text}"
        );
        assert!(matches!(&blocks[1], ContentBlock::Resource { .. }));
    }

    #[test]
    fn extract_content_with_null_structured_content_omits_it() {
        let v = json!({
            "content": [{"type": "text", "text": "ok"}],
            "structuredContent": null,
        });
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 1);
        assert_eq!(blocks[0], ContentBlock::Text { text: "ok".into() });
    }

    #[test]
    fn extract_content_without_structured_content_unchanged() {
        let v = json!({
            "content": [{"type": "text", "text": "ok"}],
            "isError": false,
        });
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 1);
        assert_eq!(blocks[0], ContentBlock::Text { text: "ok".into() });
    }

    // Strategy 4: mixed object with block-shaped field values
    #[test]
    fn extract_mixed_object_separates_blocks_and_remainder() {
        let v = json!({
            "summary": "found 3 results",
            "count": 3,
            "screenshot": {"type": "image", "mime_type": "image/png", "data": "b64"},
        });
        let blocks = extract_content_blocks(&v);
        // Remainder (summary + count) as JSON text, then the image.
        assert_eq!(blocks.len(), 2);
        // First block is the remainder text (field order in JSON objects
        // is not guaranteed, so just check it's Text and non-empty).
        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text.contains("summary")));
        assert!(matches!(&blocks[1], ContentBlock::Image { .. }));
    }

    #[test]
    fn extract_object_with_block_array_field() {
        let v = json!({
            "metadata": "info",
            "results": [
                {"type": "text", "text": "a"},
                {"type": "text", "text": "b"},
            ],
        });
        let blocks = extract_content_blocks(&v);
        // results field → 2 blocks extracted, metadata → remainder text.
        assert_eq!(blocks.len(), 3);
        assert!(matches!(&blocks[0], ContentBlock::Text { text } if text.contains("metadata")));
        assert_eq!(blocks[1], ContentBlock::Text { text: "a".into() });
        assert_eq!(blocks[2], ContentBlock::Text { text: "b".into() });
    }

    // Strategy 5: fallback
    #[test]
    fn extract_plain_string() {
        let blocks = extract_content_blocks(&json!("hello world"));
        assert_eq!(
            blocks,
            vec![ContentBlock::Text {
                text: "hello world".into()
            }]
        );
    }

    #[test]
    fn extract_number() {
        let blocks = extract_content_blocks(&json!(42));
        assert_eq!(blocks, vec![ContentBlock::Text { text: "42".into() }]);
    }

    #[test]
    fn extract_null() {
        let blocks = extract_content_blocks(&json!(null));
        assert_eq!(
            blocks,
            vec![ContentBlock::Text {
                text: "null".into()
            }]
        );
    }

    #[test]
    fn extract_plain_object_no_blocks() {
        let v = json!({"a": 1, "b": "two"});
        let blocks = extract_content_blocks(&v);
        assert_eq!(blocks.len(), 1);
        assert!(matches!(&blocks[0], ContentBlock::Text { .. }));
    }

    /// `extract_content_blocks` surfaces a `ToolRunResult`'s `prompt_text` as the
    /// model content, not a JSON dump of the struct.
    #[test]
    fn tool_run_result_shape_extracts_prompt_text() {
        let prompt = "1: a.txt\n2: b.txt\n<system-reminder>\nThe todo_write tool \
                      hasn't been used recently.\n</system-reminder>";
        let v = json!({
            "output": {"list_dir": {"entries": ["a.txt", "b.txt"]}},
            "prompt_text": prompt,
            "effective_tool_name": null,
        });
        assert_eq!(
            extract_content_blocks(&v),
            vec![ContentBlock::Text {
                text: prompt.to_owned()
            }]
        );
    }
}