orchestral-runtime 0.2.0

mod catalog;
mod http;
mod parsing;
mod prompt;

use std::collections::BTreeSet;
use std::sync::Arc;

use async_trait::async_trait;
use thiserror::Error;
use tracing::{debug, info, warn};

use orchestral_core::planner::{PlanError, Planner, PlannerContext, PlannerOutput};
use orchestral_core::types::Intent;

pub use self::http::{HttpLlmClient, HttpLlmClientConfig};
use self::parsing::{extract_json, parse_action_selection, parse_planner_output, ActionSelection};
use self::prompt::{build_action_selector_prompt, build_planner_prompt, truncate_for_log};

const MAX_PROMPT_LOG_CHARS: usize = 4_000;
const MAX_LLM_OUTPUT_LOG_CHARS: usize = 8_000;
/// LLM request payload
#[derive(Debug, Clone)]
pub struct LlmRequest {
    pub system: String,
    pub user: String,
    pub model: String,
    pub temperature: f32,
}

/// Tool definition for LLM function calling.
#[derive(Debug, Clone)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    /// JSON Schema describing the tool parameters.
    pub parameters: serde_json::Value,
}

/// LLM response that may contain a structured tool call instead of text.
#[derive(Debug, Clone)]
pub enum LlmResponse {
    Text(String),
    ToolCall {
        id: String,
        name: String,
        arguments: serde_json::Value,
    },
}

pub type StreamChunkCallback = Arc<dyn Fn(String) + Send + Sync>;

/// LLM client trait
#[async_trait]
pub trait LlmClient: Send + Sync {
    async fn complete(&self, request: LlmRequest) -> Result<String, LlmError>;

    /// Chat with tool definitions. LLM must call one of the provided tools.
    /// Default implementation falls back to text-only `complete`.
    async fn complete_with_tools(
        &self,
        request: LlmRequest,
        tools: &[ToolDefinition],
    ) -> Result<LlmResponse, LlmError> {
        let _ = tools;
        let text = self.complete(request).await?;
        Ok(LlmResponse::Text(text))
    }

    async fn complete_stream(
        &self,
        request: LlmRequest,
        on_chunk: StreamChunkCallback,
    ) -> Result<String, LlmError> {
        let full = self.complete(request).await?;
        for token in full.split_inclusive(char::is_whitespace) {
            if !token.is_empty() {
                on_chunk(token.to_string());
            }
        }
        Ok(full)
    }
}

#[async_trait]
impl LlmClient for Arc<dyn LlmClient> {
    async fn complete(&self, request: LlmRequest) -> Result<String, LlmError> {
        (**self).complete(request).await
    }

    async fn complete_with_tools(
        &self,
        request: LlmRequest,
        tools: &[ToolDefinition],
    ) -> Result<LlmResponse, LlmError> {
        (**self).complete_with_tools(request, tools).await
    }

    async fn complete_stream(
        &self,
        request: LlmRequest,
        on_chunk: StreamChunkCallback,
    ) -> Result<String, LlmError> {
        (**self).complete_stream(request, on_chunk).await
    }
}

/// LLM errors
#[derive(Debug, Error)]
pub enum LlmError {
    #[error("http error: {0}")]
    Http(String),
    #[error("response error: {0}")]
    Response(String),
    #[error("serialization error: {0}")]
    Serialization(String),
}

/// Planner config for LLM
#[derive(Debug, Clone)]
pub struct LlmPlannerConfig {
    pub model: String,
    pub temperature: f32,
    pub max_history: usize,
    pub system_prompt: String,
    pub log_full_prompts: bool,
    pub selector_min_action_count: usize,
    pub selector_max_actions: usize,
}

impl Default for LlmPlannerConfig {
    fn default() -> Self {
        Self {
            model: "anthropic/claude-sonnet-4.5".to_string(),
            temperature: 0.2,
            max_history: 20,
            system_prompt: String::new(),
            log_full_prompts: false,
            selector_min_action_count: 30,
            selector_max_actions: 30,
        }
    }
}

/// LLM-based planner
pub struct LlmPlanner<C: LlmClient> {
    pub client: C,
    pub config: LlmPlannerConfig,
}

impl<C: LlmClient> LlmPlanner<C> {
    pub fn new(client: C, config: LlmPlannerConfig) -> Self {
        Self { client, config }
    }

    fn build_prompt(&self, intent: &Intent, context: &PlannerContext) -> (String, String) {
        build_planner_prompt(
            &self.config.system_prompt,
            intent,
            context,
            self.config.max_history,
        )
    }

    fn build_selector_prompt(&self, intent: &Intent, context: &PlannerContext) -> (String, String) {
        build_action_selector_prompt(
            &self.config.system_prompt,
            intent,
            context,
            self.config.max_history,
            self.config.selector_max_actions,
        )
    }

    fn should_run_action_selector(&self, context: &PlannerContext) -> bool {
        let action_count = context.available_actions.len();
        action_count >= self.config.selector_min_action_count
            && action_count > self.config.selector_max_actions
    }

    fn apply_selection(
        &self,
        context: &PlannerContext,
        selection: ActionSelection,
    ) -> Option<ResolvedActionSelection> {
        let selected = selection
            .selected_actions
            .into_iter()
            .collect::<BTreeSet<_>>();
        let blocked = selection
            .blocked_actions
            .into_iter()
            .collect::<BTreeSet<_>>();

        let mut resolved_actions = Vec::new();
        let mut resolved_selected_names = Vec::new();
        let mut resolved_blocked_names = Vec::new();

        for action in &context.available_actions {
            if blocked.contains(&action.name) {
                resolved_blocked_names.push(action.name.clone());
                continue;
            }
            if selected.contains(&action.name)
                && resolved_actions.len() < self.config.selector_max_actions
            {
                resolved_selected_names.push(action.name.clone());
                resolved_actions.push(action.clone());
            }
        }

        if resolved_actions.is_empty() {
            return None;
        }

        if resolved_actions.len() >= context.available_actions.len()
            && resolved_blocked_names.is_empty()
        {
            return None;
        }

        let filtered_context = PlannerContext {
            available_actions: resolved_actions,
            history: context.history.clone(),
            runtime_info: context.runtime_info.clone(),
            skill_instructions: context.skill_instructions.clone(),
            skill_summaries: context.skill_summaries.clone(),
            loop_context: context.loop_context.clone(),
        };

        Some(ResolvedActionSelection {
            filtered_context,
            selected_actions: resolved_selected_names,
            blocked_actions: resolved_blocked_names,
            reason: selection.reason,
        })
    }

    async fn maybe_select_actions(
        &self,
        intent: &Intent,
        context: &PlannerContext,
    ) -> Result<Option<ResolvedActionSelection>, PlanError> {
        if !self.should_run_action_selector(context) {
            return Ok(None);
        }

        let (system, user) = self.build_selector_prompt(intent, context);
        info!(
            model = %self.config.model,
            temperature = self.config.temperature,
            action_count = context.available_actions.len(),
            selector_max_actions = self.config.selector_max_actions,
            selector_min_action_count = self.config.selector_min_action_count,
            "action selector request prepared"
        );
        if tracing::enabled!(tracing::Level::DEBUG) {
            if self.config.log_full_prompts {
                debug!(
                    system_prompt = %system,
                    user_prompt = %user,
                    system_chars = system.chars().count(),
                    user_chars = user.chars().count(),
                    "action selector prompts (full)"
                );
            } else {
                debug!(
                    system_prompt = %truncate_for_log(&system, MAX_PROMPT_LOG_CHARS),
                    user_prompt = %truncate_for_log(&user, MAX_PROMPT_LOG_CHARS),
                    system_chars = system.chars().count(),
                    user_chars = user.chars().count(),
                    "action selector prompts"
                );
            }
        }

        let request = LlmRequest {
            system,
            user,
            model: self.config.model.clone(),
            temperature: self.config.temperature,
        };
        let output = self
            .client
            .complete(request)
            .await
            .map_err(|e| PlanError::LlmError(e.to_string()))?;
        if tracing::enabled!(tracing::Level::DEBUG) {
            debug!(
                llm_output = %truncate_for_log(&output, MAX_LLM_OUTPUT_LOG_CHARS),
                "action selector raw llm output"
            );
        }

        let json_str = match extract_json(&output) {
            Some(json) => json,
            None => {
                warn!(
                    "action selector output did not contain JSON; falling back to full action set"
                );
                return Ok(None);
            }
        };
        if tracing::enabled!(tracing::Level::DEBUG) {
            debug!(
                selector_json = %truncate_for_log(&json_str, MAX_LLM_OUTPUT_LOG_CHARS),
                "action selector extracted json"
            );
        }

        let selection = match parse_action_selection(&json_str) {
            Ok(selection) => selection,
            Err(error) => {
                warn!(error = %error, "action selector output was invalid; falling back to full action set");
                return Ok(None);
            }
        };

        let resolved = match self.apply_selection(context, selection) {
            Some(resolved) => resolved,
            None => {
                warn!("action selector did not resolve any narrower action subset; falling back to full action set");
                return Ok(None);
            }
        };

        info!(
            selected_count = resolved.selected_actions.len(),
            blocked_count = resolved.blocked_actions.len(),
            selected_actions = %resolved.selected_actions.join(", "),
            blocked_actions = %resolved.blocked_actions.join(", "),
            reason = ?resolved.reason,
            "action selector resolved actions"
        );
        Ok(Some(resolved))
    }
}

struct ResolvedActionSelection {
    filtered_context: PlannerContext,
    selected_actions: Vec<String>,
    blocked_actions: Vec<String>,
    reason: Option<String>,
}

#[async_trait]
impl<C: LlmClient> Planner for LlmPlanner<C> {
    async fn plan(
        &self,
        intent: &Intent,
        context: &PlannerContext,
    ) -> Result<PlannerOutput, PlanError> {
        let selected_context = self
            .maybe_select_actions(intent, context)
            .await?
            .map(|selection| selection.filtered_context);
        let planner_context = selected_context.as_ref().unwrap_or(context);

        let (system, user) = self.build_prompt(intent, planner_context);
        info!(
            model = %self.config.model,
            temperature = self.config.temperature,
            intent_len = intent.content.len(),
            action_count = planner_context.available_actions.len(),
            history_count = planner_context.history.len(),
            "planner request prepared"
        );
        if tracing::enabled!(tracing::Level::DEBUG) {
            if self.config.log_full_prompts {
                debug!(
                    system_prompt = %system,
                    user_prompt = %user,
                    system_chars = system.chars().count(),
                    user_chars = user.chars().count(),
                    "planner prompts (full)"
                );
            } else {
                let system_preview = truncate_for_log(&system, MAX_PROMPT_LOG_CHARS);
                let user_preview = truncate_for_log(&user, MAX_PROMPT_LOG_CHARS);
                debug!(
                    system_prompt = %system_preview,
                    user_prompt = %user_preview,
                    system_chars = system.chars().count(),
                    user_chars = user.chars().count(),
                    "planner prompts"
                );
            }
        }
        let request = LlmRequest {
            system,
            user,
            model: self.config.model.clone(),
            temperature: self.config.temperature,
        };
        let output = self
            .client
            .complete(request)
            .await
            .map_err(|e| PlanError::LlmError(e.to_string()))?;
        if tracing::enabled!(tracing::Level::DEBUG) {
            debug!(
                llm_output = %truncate_for_log(&output, MAX_LLM_OUTPUT_LOG_CHARS),
                "planner raw llm output"
            );
        }

        let json_str = extract_json(&output)
            .ok_or_else(|| PlanError::Generation("LLM output did not contain JSON".to_string()))?;
        if tracing::enabled!(tracing::Level::DEBUG) {
            debug!(
                plan_json = %truncate_for_log(&json_str, MAX_LLM_OUTPUT_LOG_CHARS),
                "planner extracted json"
            );
        }

        let output = parse_planner_output(&json_str)?;
        match &output {
            PlannerOutput::SingleAction(call) => {
                info!(
                    output_type = "single_action",
                    action = %call.action,
                    reason = ?call.reason,
                    "planner parsed output"
                );
            }
            PlannerOutput::MiniPlan(plan) => {
                info!(
                    output_type = "mini_plan",
                    goal = %plan.goal,
                    step_count = plan.steps.len(),
                    "planner parsed output"
                );
            }
            PlannerOutput::Done(message) => {
                info!(
                    output_type = "done",
                    message = %truncate_for_log(message, MAX_PROMPT_LOG_CHARS),
                    "planner parsed output"
                );
            }
            PlannerOutput::NeedInput(question) => {
                info!(
                    output_type = "need_input",
                    question = %truncate_for_log(question, MAX_PROMPT_LOG_CHARS),
                    "planner parsed output"
                );
            }
        }
        Ok(output)
    }
}

/// Mock LLM client for tests/examples
pub struct MockLlmClient {
    pub response: String,
}

#[async_trait]
impl LlmClient for MockLlmClient {
    async fn complete(&self, _request: LlmRequest) -> Result<String, LlmError> {
        Ok(self.response.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use orchestral_core::action::ActionMeta;
    use orchestral_core::planner::{PlannerContext, SkillInstruction};
    use orchestral_core::types::Intent;
    use serde_json::json;
    use std::collections::VecDeque;
    use std::sync::Mutex;

    #[test]
    fn test_planner_prompt_uses_new_output_shapes() {
        let planner = LlmPlanner::new(
            MockLlmClient {
                response: "{}".to_string(),
            },
            LlmPlannerConfig {
                system_prompt: "Base prompt.".to_string(),
                ..LlmPlannerConfig::default()
            },
        );

        let actions = vec![
            ActionMeta::new("write_doc", "Write markdown to file")
                .with_capabilities(["filesystem_write", "side_effect"])
                .with_input_kinds(["path", "text"])
                .with_output_kinds(["path"])
                .with_input_schema(json!({
                    "type":"object",
                    "properties":{
                        "path":{"type":"string","description":"Target markdown path","example":"guide.md"},
                        "content":{"type":"string","description":"Markdown content"}
                    },
                    "required":["path","content"]
                }))
                .with_output_schema(
                    json!({
                        "type":"object",
                        "properties":{
                            "path":{"type":"string","description":"Resolved path"},
                            "bytes":{"type":"integer","description":"Written bytes"}
                        }
                    }),
                ),
            ActionMeta::new("file_read", "Read a file")
                .with_capabilities(["filesystem_read"])
                .with_input_kinds(["path"])
                .with_output_kinds(["text"]),
        ];
        let context = PlannerContext::new(actions);
        let intent = Intent::new("generate a guide");
        let (system, user) = planner.build_prompt(&intent, &context);

        assert!(system.contains("Orchestral Planner"));
        assert!(system.contains("Legacy workflow/stage outputs are disabled."));
        assert!(system.contains("SINGLE_ACTION"));
        assert!(system.contains("MINI_PLAN"));
        assert!(!system.contains("Action Catalog"));
        assert!(user.contains("\"type\":\"SINGLE_ACTION\""));
        assert!(user.contains("\"type\":\"MINI_PLAN\""));
        assert!(user.contains("\"type\":\"DONE\""));
        assert!(user.contains("\"type\":\"NEED_INPUT\""));
        assert!(!user.contains("\"type\":\"WORKFLOW\""));
        assert!(!user.contains("\"type\":\"STAGE_CHOICE\""));
        assert!(user.contains("DONE must never claim to execute commands"));
    }

    #[test]
    fn test_planner_prompt_contains_skill_knowledge() {
        let planner = LlmPlanner::new(
            MockLlmClient {
                response: "{}".to_string(),
            },
            LlmPlannerConfig::default(),
        );

        let actions = vec![ActionMeta::new("mcp__alpha", "Call MCP server alpha")
            .with_capabilities(["mcp", "side_effect"])
            .with_input_kinds(["structured"])
            .with_output_kinds(["structured"])
            .with_input_schema(json!({
                "type":"object",
                "properties":{
                    "operation":{"type":"string"},
                    "tool":{"type":"string"},
                    "arguments":{"type":"object"}
                }
            }))
            .with_output_schema(json!({
                "type":"object",
                "properties":{
                    "server":{"type":"string"},
                    "result":{}
                }
            }))];
        let context =
            PlannerContext::new(actions).with_skill_instructions(vec![SkillInstruction {
                skill_name: "demo".to_string(),
                instructions: "Always write then verify.".to_string(),
                skill_path: Some("skills/demo/SKILL.md".to_string()),
                scripts_dir: Some(".claude/skills/demo/scripts".to_string()),
                venv_python: None,
            }]);
        let intent = Intent::new("need tools and skills");
        let (system, _user) = planner.build_prompt(&intent, &context);

        assert!(system.contains("Activated Skills:"));
        assert!(system.contains("never invent script filenames"));
        assert!(system.contains("- demo"));
        assert!(system.contains("Always write then verify."));
        assert!(system.contains("[skill file: skills/demo/SKILL.md]"));
        assert!(system.contains("[scripts: .claude/skills/demo/scripts]"));
        assert!(!system.contains("Action Catalog"));
    }

    #[test]
    fn test_parse_done_output() {
        let raw = r#"{"type":"DONE","message":"你好"}"#;
        let parsed = parse_planner_output(raw).expect("parse done");
        match parsed {
            PlannerOutput::Done(message) => {
                assert_eq!(message, "你好");
            }
            _ => panic!("expected done output"),
        }
    }

    #[test]
    fn test_parse_single_action_output() {
        let raw = r#"{"type":"SINGLE_ACTION","action":"file_read","params":{"path":"README.md"},"reason":"read readme"}"#;
        let parsed = parse_planner_output(raw).expect("parse single action");
        match parsed {
            PlannerOutput::SingleAction(call) => {
                assert_eq!(call.action, "file_read");
                assert_eq!(call.params["path"], "README.md");
                assert_eq!(call.reason.as_deref(), Some("read readme"));
            }
            _ => panic!("expected single_action output"),
        }
    }

    #[test]
    fn test_parse_need_input_output() {
        let raw = r#"{"type":"NEED_INPUT","question":"请提供文件路径"}"#;
        let parsed = parse_planner_output(raw).expect("parse need_input");
        match parsed {
            PlannerOutput::NeedInput(question) => {
                assert_eq!(question, "请提供文件路径");
            }
            _ => panic!("expected need_input output"),
        }
    }

    #[test]
    fn test_parse_mini_plan_output() {
        let raw = r#"{
            "type":"MINI_PLAN",
            "goal":"inspect then summarize",
            "steps":[
                {"id":"list_docs","action":"shell","params":{"command":"find ./docs -maxdepth 1 -type f"}},
                {"id":"read_readme","action":"file_read","depends_on":["list_docs"],"params":{"path":"README.md"}}
            ],
            "on_complete":"{{read_readme.content}}"
        }"#;
        let parsed = parse_planner_output(raw).expect("parse mini plan");
        match parsed {
            PlannerOutput::MiniPlan(plan) => {
                assert_eq!(plan.goal, "inspect then summarize");
                assert_eq!(plan.steps.len(), 2);
                assert_eq!(plan.steps[1].depends_on.len(), 1);
                assert_eq!(plan.on_complete.as_deref(), Some("{{read_readme.content}}"));
            }
            _ => panic!("expected mini_plan output"),
        }
    }

    #[test]
    fn test_extract_json_ignores_non_json_braces() {
        let raw = r#"Preface {not json} -> {"type":"DONE","message":"ok"} trailing"#;
        let json = extract_json(raw).expect("json");
        assert_eq!(json, r#"{"type":"DONE","message":"ok"}"#);
    }

    #[test]
    fn test_extract_json_handles_braces_inside_strings() {
        let raw = r#"noise {"type":"DONE","message":"value with } brace"} end"#;
        let json = extract_json(raw).expect("json");
        assert_eq!(json, r#"{"type":"DONE","message":"value with } brace"}"#);
    }

    #[test]
    fn test_parse_action_selection_output() {
        let raw = r#"{"selected_actions":["file_read","file_write"],"blocked_actions":["shell"],"reason":"typed actions are enough"}"#;
        let parsed = parse_action_selection(raw).expect("parse action selection");

        assert_eq!(parsed.selected_actions, vec!["file_read", "file_write"]);
        assert_eq!(parsed.blocked_actions, vec!["shell"]);
        assert_eq!(parsed.reason.as_deref(), Some("typed actions are enough"));
    }

    struct RecordingMockLlmClient {
        responses: Mutex<VecDeque<String>>,
        requests: Mutex<Vec<LlmRequest>>,
    }

    #[async_trait]
    impl LlmClient for RecordingMockLlmClient {
        async fn complete(&self, request: LlmRequest) -> Result<String, LlmError> {
            self.requests.lock().expect("requests lock").push(request);
            self.responses
                .lock()
                .expect("responses lock")
                .pop_front()
                .ok_or_else(|| LlmError::Response("no queued mock response".to_string()))
        }
    }

    #[tokio::test]
    async fn test_llm_planner_uses_selector_filtered_actions() {
        let planner = LlmPlanner::new(
            RecordingMockLlmClient {
                responses: Mutex::new(VecDeque::from(vec![
                    r#"{"selected_actions":["file_read","file_write"],"blocked_actions":["shell"],"reason":"typed file actions are sufficient"}"#.to_string(),
                    r#"{"type":"DONE","message":"ok"}"#.to_string(),
                ])),
                requests: Mutex::new(Vec::new()),
            },
            LlmPlannerConfig {
                selector_min_action_count: 1,
                selector_max_actions: 2,
                ..LlmPlannerConfig::default()
            },
        );

        let context = PlannerContext::new(vec![
            ActionMeta::new("shell", "shell"),
            ActionMeta::new("file_read", "read"),
            ActionMeta::new("file_write", "write"),
        ]);
        let result = planner
            .plan(&Intent::new("read and write the file safely"), &context)
            .await
            .expect("planner result");

        match result {
            PlannerOutput::Done(message) => assert_eq!(message, "ok"),
            other => panic!("expected done output, got {:?}", other),
        }

        let requests = planner.client.requests.lock().expect("requests lock");
        assert_eq!(requests.len(), 2);
        assert!(requests[0].system.contains("Orchestral Action Selector."));
        assert!(requests[1].system.contains("- file_read: read"));
        assert!(requests[1].system.contains("- file_write: write"));
        assert!(!requests[1].system.contains("- shell: shell"));
    }

    #[tokio::test]
    async fn test_llm_planner_falls_back_when_selector_returns_unknown_actions() {
        let planner = LlmPlanner::new(
            RecordingMockLlmClient {
                responses: Mutex::new(VecDeque::from(vec![
                    r#"{"selected_actions":["unknown_action"],"blocked_actions":[],"reason":"bad output"}"#.to_string(),
                    r#"{"type":"DONE","message":"ok"}"#.to_string(),
                ])),
                requests: Mutex::new(Vec::new()),
            },
            LlmPlannerConfig {
                selector_min_action_count: 1,
                selector_max_actions: 1,
                ..LlmPlannerConfig::default()
            },
        );

        planner
            .plan(
                &Intent::new("inspect the workspace"),
                &PlannerContext::new(vec![
                    ActionMeta::new("shell", "shell"),
                    ActionMeta::new("file_read", "read"),
                ]),
            )
            .await
            .expect("planner result");

        let requests = planner.client.requests.lock().expect("requests lock");
        assert_eq!(requests.len(), 2);
        assert!(requests[1].system.contains("- shell: shell"));
        assert!(requests[1].system.contains("- file_read: read"));
    }
}