//! Sub-agent session loading.
//!
//! Claude Code spawns sub-agents for Agent tool calls and stores their
//! conversations under `<session-artifacts-dir>/subagents/`. Each sub-agent
//! has a JSONL file and an optional `.meta.json` with metadata such as the
//! agent type.

use camino::Utf8Path;

use crate::error::Result;
use crate::parser::parse_session;
use crate::types::{SubAgentMeta, SubAgentSession, ToolExchange};

/// Load all sub-agent sessions from a session's artifacts directory.
///
/// Expects files at: `<session_artifacts_dir>/subagents/agent-<id>.jsonl`
/// with optional companion: `<session_artifacts_dir>/subagents/agent-<id>.meta.json`
///
/// Returns an empty `Vec` if the `subagents/` subdirectory does not exist.
pub fn load_subagents(session_artifacts_dir: &Utf8Path) -> Result<Vec<SubAgentSession>> {
    let subagents_dir = session_artifacts_dir.join("subagents");
    if !subagents_dir.as_std_path().exists() {
        return Ok(Vec::new());
    }

    let mut subagents = Vec::new();

    for entry in std::fs::read_dir(subagents_dir.as_std_path())? {
        let entry = entry?;
        let path = entry.path();

        // Only process agent-*.jsonl files
        let file_name = match path.file_name().and_then(|n| n.to_str()) {
            Some(name) => name.to_string(),
            None => continue,
        };

        if !file_name.starts_with("agent-") || !file_name.ends_with(".jsonl") {
            continue;
        }

        // Extract agent ID: "agent-abc123.jsonl" → "abc123"
        let agent_id = file_name
            .strip_prefix("agent-")
            .and_then(|s| s.strip_suffix(".jsonl"))
            .unwrap_or(&file_name)
            .to_string();

        let jsonl_path = match camino::Utf8PathBuf::try_from(path.clone()) {
            Ok(p) => p,
            Err(_) => continue,
        };

        // Parse the JSONL session
        let messages = parse_session(&jsonl_path)?;

        // Try to load companion metadata file
        let meta_path = path.with_file_name(format!("agent-{}.meta.json", agent_id));
        let meta = if meta_path.exists() {
            let content = std::fs::read_to_string(&meta_path)?;
            serde_json::from_str::<SubAgentMeta>(&content)
                .unwrap_or(SubAgentMeta { agent_type: None })
        } else {
            SubAgentMeta { agent_type: None }
        };

        subagents.push(SubAgentSession {
            agent_id,
            meta,
            messages,
            parent_tool_call_id: None,
        });
    }

    Ok(subagents)
}

/// Try to link sub-agent sessions to their parent Agent tool calls.
///
/// For each Agent tool call in `tool_exchanges`, parses the tool result text
/// for `agentId: <id>` and matches it against sub-agent sessions. On match,
/// sets `SubAgentSession.parent_tool_call_id` to the tool call's ID.
///
/// This is best-effort: if the agentId text format changes or a result is
/// missing, the sub-agent is still usable but without tool_use linkage.
pub fn link_subagents_to_calls(
    subagents: &mut [SubAgentSession],
    tool_exchanges: &[ToolExchange],
) {
    use regex::Regex;

    if subagents.is_empty() || tool_exchanges.is_empty() {
        return;
    }

    // Compile once, match many
    let re = Regex::new(r"agentId:\s*(\S+)").expect("valid regex");

    for exchange in tool_exchanges {
        // Only look at Agent tool calls
        if exchange.call.name != crate::types::ToolName::Agent {
            continue;
        }

        // Extract agentId from the tool result text
        let agent_id = exchange
            .result
            .as_ref()
            .and_then(|r| r.content.as_deref())
            .and_then(|text| re.captures(text))
            .and_then(|caps| caps.get(1))
            .map(|m| m.as_str());

        let agent_id = match agent_id {
            Some(id) => id,
            None => continue,
        };

        // Find matching sub-agent and set the linkage
        if let Some(subagent) = subagents.iter_mut().find(|s| s.agent_id == agent_id) {
            subagent.parent_tool_call_id = Some(exchange.call.id.clone());
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{ToolCall, ToolName, ToolResultData};

    #[test]
    fn test_link_subagents_to_calls_matches_agent_id() {
        let mut subagents = vec![
            SubAgentSession {
                agent_id: "abc123def".to_string(),
                meta: SubAgentMeta { agent_type: Some("Explore".to_string()) },
                messages: vec![],
                parent_tool_call_id: None,
            },
            SubAgentSession {
                agent_id: "xyz789".to_string(),
                meta: SubAgentMeta { agent_type: None },
                messages: vec![],
                parent_tool_call_id: None,
            },
        ];

        let exchanges = vec![
            ToolExchange {
                call: ToolCall {
                    id: "toolu_01ABC".to_string(),
                    name: ToolName::Agent,
                    input: serde_json::json!({"description": "test"}),
                    source_message_uuid: "msg-1".to_string(),
                },
                result: Some(ToolResultData {
                    tool_use_id: "toolu_01ABC".to_string(),
                    content: Some("agentId: abc123def (use SendMessage with to: 'abc123def' to continue)\n<usage>total_tokens: 1000</usage>".to_string()),
                    is_error: false,
                    source_message_uuid: "msg-2".to_string(),
                }),
            },
            ToolExchange {
                call: ToolCall {
                    id: "toolu_02DEF".to_string(),
                    name: ToolName::Read,
                    input: serde_json::json!({}),
                    source_message_uuid: "msg-3".to_string(),
                },
                result: None,
            },
        ];

        link_subagents_to_calls(&mut subagents, &exchanges);

        assert_eq!(subagents[0].parent_tool_call_id, Some("toolu_01ABC".to_string()));
        assert_eq!(subagents[1].parent_tool_call_id, None);
    }

    #[test]
    fn test_link_subagents_empty_inputs() {
        let mut empty_subagents: Vec<SubAgentSession> = vec![];
        let empty_exchanges: Vec<ToolExchange> = vec![];
        link_subagents_to_calls(&mut empty_subagents, &empty_exchanges);
        // No panic
    }

    #[test]
    fn test_link_subagents_no_match() {
        let mut subagents = vec![SubAgentSession {
            agent_id: "no_match".to_string(),
            meta: SubAgentMeta { agent_type: None },
            messages: vec![],
            parent_tool_call_id: None,
        }];

        let exchanges = vec![ToolExchange {
            call: ToolCall {
                id: "toolu_99".to_string(),
                name: ToolName::Agent,
                input: serde_json::json!({}),
                source_message_uuid: "msg-1".to_string(),
            },
            result: Some(ToolResultData {
                tool_use_id: "toolu_99".to_string(),
                content: Some("agentId: different_id\n<usage>tokens: 500</usage>".to_string()),
                is_error: false,
                source_message_uuid: "msg-2".to_string(),
            }),
        }];

        link_subagents_to_calls(&mut subagents, &exchanges);
        assert_eq!(subagents[0].parent_tool_call_id, None);
    }
}
