//! AgentRunner — executes agent tasks through the turn loop.
pub(crate) mod compaction;
pub(crate) mod turn_loop;
pub mod types;

use std::sync::Arc;

use anyhow::Result;
use tokio::sync::mpsc;

use nenjo_models::ChatMessage;
use tracing::{debug, info, trace};
use uuid::Uuid;

use super::abilities::{build_ability_tools, is_ability_tool};
use super::instance::AgentExecutionMode;
use super::sub_agents::{
    ChildRuntimeHandle, PARENT_TOOL_NAMES, SubAgentLimits, SubAgentRuntime, child_tools,
    parent_tools,
};
use anyhow::Context;
use nenjo_models::ModelProvider;

use super::instance::AgentInstance;
use crate::Slug;
use crate::input::{AgentRun, AgentRunKind, ChatInput, TaskInput};
use crate::manifest::{AbilityManifest, DomainManifest, Manifest};
use crate::memory::{self, MemoryScope};
use crate::provider::{ErasedProvider, ProviderRuntime};
use crate::types::ActiveDomain;
use types::{TurnEvent, TurnOutput};

/// Handle to a running agent execution.
///
/// Provides a stream of [`TurnEvent`]s as the agent works, plus access
/// to the final [`TurnOutput`] when done.
pub struct ExecutionHandle {
    events_rx: mpsc::UnboundedReceiver<TurnEvent>,
    join: Option<tokio::task::JoinHandle<Result<TurnOutput>>>,
    pause_token: types::PauseToken,
}

impl ExecutionHandle {
    /// Receive the next event. Returns `None` when the turn loop finishes.
    pub async fn recv(&mut self) -> Option<TurnEvent> {
        self.events_rx.recv().await
    }

    /// Get a mutable reference to the underlying event receiver.
    pub fn events(&mut self) -> &mut mpsc::UnboundedReceiver<TurnEvent> {
        &mut self.events_rx
    }

    /// Get a clone of the pause token for external control (e.g. execution registry).
    pub fn pause_token(&self) -> types::PauseToken {
        self.pause_token.clone()
    }

    /// Abort the running execution. The spawned task is cancelled immediately.
    pub fn abort(&self) {
        if let Some(join) = &self.join {
            join.abort();
        }
    }

    /// Pause execution. The turn loop will block before the next LLM call.
    ///
    /// The caller will receive a `TurnEvent::Paused` event once the loop
    /// reaches the pause point. In-flight tool executions finish first.
    pub fn pause(&self) {
        self.pause_token.pause();
    }

    /// Resume a paused execution. The turn loop continues from where it stopped.
    ///
    /// The caller will receive a `TurnEvent::Resumed` event.
    pub fn resume(&self) {
        self.pause_token.resume();
    }

    /// Check if the execution is currently paused.
    pub fn is_paused(&self) -> bool {
        self.pause_token.is_paused()
    }

    /// Wait for the final output.
    pub async fn output(mut self) -> Result<TurnOutput> {
        self.join
            .take()
            .expect("execution join handle should be present")
            .await
            .map_err(|e| anyhow::anyhow!("execution task panicked: {e}"))?
    }
}

impl Drop for ExecutionHandle {
    fn drop(&mut self) {
        if let Some(join) = &self.join
            && !join.is_finished()
        {
            join.abort();
        }
    }
}

pub(crate) fn build_instruction_messages(
    system_prompt: &str,
    developer_prompt: &str,
    supports_developer_role: bool,
) -> Vec<ChatMessage> {
    let combined = [system_prompt, developer_prompt]
        .into_iter()
        .filter(|part| !part.trim().is_empty())
        .collect::<Vec<_>>()
        .join("\n\n");

    if combined.is_empty() {
        Vec::new()
    } else if supports_developer_role {
        vec![ChatMessage::developer(combined)]
    } else {
        vec![ChatMessage::system(combined)]
    }
}

/// Wraps an [`AgentInstance`] and provides the execution API.
///
/// Created via [`AgentBuilder::build()`](super::builder::AgentBuilder::build).
pub struct AgentRunner<P: ProviderRuntime = ErasedProvider> {
    instance: Arc<AgentInstance<P>>,
    memory: Option<Arc<P::Memory<'static>>>,
    memory_scope: Option<MemoryScope>,
    manifest: Option<Arc<Manifest>>,
}

impl<P: ProviderRuntime> AgentRunner<P> {
    pub(crate) async fn new(
        mut instance: AgentInstance<P>,
        memory: Option<Arc<P::Memory<'static>>>,
        memory_scope: Option<MemoryScope>,
    ) -> Result<Self, super::error::AgentError> {
        // Extract manifest before delegation is consumed so domain_expansion can
        // pass it to sub-runners.
        let manifest = instance
            .runtime
            .provider_runtime
            .as_ref()
            .map(|provider| provider.manifest_snapshot());

        // If the agent has abilities, register the ability discovery and
        // invocation broker tools resolved from the canonical manifest.
        if let Some(active_abilities) = instance
            .runtime
            .provider_runtime
            .as_ref()
            .filter(|_| instance.runtime.execution_mode == AgentExecutionMode::Parent)
            .map(|provider| resolve_active_abilities(provider, &instance.manifest.abilities, None))
            .filter(|abilities| !abilities.is_empty())
        {
            let base_instance = Arc::new(instance.clone());
            instance
                .runtime
                .tools
                .extend(build_ability_tools(&active_abilities, base_instance)?);
        }

        let instance = Arc::new(instance);

        Ok(Self {
            instance,
            memory,
            memory_scope,
            manifest,
        })
    }

    /// Read-only access to the underlying agent instance.
    pub fn instance(&self) -> &AgentInstance<P> {
        &self.instance
    }

    /// The agent's name.
    pub fn agent_name(&self) -> &str {
        self.instance.name()
    }

    /// The agent's manifest ID.
    pub fn agent_id(&self) -> Uuid {
        self.instance.agent_id()
    }

    /// Create a runner from a pre-built instance.
    ///
    /// Used by the harness to re-use a domain-expanded instance across
    /// multiple chat turns without rebuilding from the Provider each time.
    /// Pass memory/scope to preserve memory and artifact loading.
    pub fn from_instance(
        instance: AgentInstance<P>,
        memory: Option<Arc<P::Memory<'static>>>,
        memory_scope: Option<MemoryScope>,
    ) -> Self {
        Self {
            instance: Arc::new(instance),
            memory,
            memory_scope,
            manifest: None,
        }
    }

    /// The memory backend, if configured.
    pub fn memory(&self) -> Option<&Arc<P::Memory<'static>>> {
        self.memory.as_ref()
    }

    /// The memory scope, if configured.
    pub fn memory_scope(&self) -> Option<&MemoryScope> {
        self.memory_scope.as_ref()
    }

    /// Activate a domain by name, returning a new runner with expanded config.
    ///
    /// The domain is looked up from the agent's assigned domains. The returned
    /// runner appends the domain's prompt addon and layers in any domain-scoped
    /// ability and MCP activations.
    ///
    /// ```ignore
    /// let domain_runner = runner.domain_expansion("prd")?;
    /// let output = domain_runner.chat("Create a PRD for auth").await?;
    /// ```
    pub async fn domain_expansion(&self, domain_name: &str) -> Result<AgentRunner<P>> {
        let provider = self
            .instance
            .runtime
            .provider_runtime
            .as_ref()
            .ok_or_else(|| {
                super::error::AgentError::MissingManifest(self.instance.name().into())
            })?;
        let domain = provider
            .find_domain(domain_name)
            .filter(|domain| domain_is_assigned(&self.instance.manifest.domains, domain))
            .with_context(|| {
                let available: Vec<&str> = self
                    .instance
                    .manifest
                    .domains
                    .iter()
                    .filter_map(|domain_slug| provider.find_domain(domain_slug.as_str()))
                    .map(|domain| domain.name.as_str())
                    .collect();
                format!("domain '{domain_name}' not found. Available: {available:?}")
            })?;

        let session_manifest: DomainManifest = domain.clone();

        // Build the active domain session state.
        let active_domain = ActiveDomain {
            session_id: Uuid::new_v4(),
            domain_id: domain.id,
            domain_name: domain.name.clone(),
            manifest: session_manifest.clone(),
        };

        // Clone the instance and apply domain expansion.
        let mut instance = (*self.instance).clone();
        instance.prompt.context.active_domain = Some(active_domain);

        let active_abilities = resolve_active_abilities(
            provider,
            &instance.manifest.abilities,
            Some(&session_manifest),
        );

        // Rebuild ability broker tools so the visible set matches the
        // effective domain-expanded ability assignments.
        instance
            .runtime
            .tools
            .retain(|tool| !is_ability_tool(tool.name()));
        if !active_abilities.is_empty() {
            let base_instance = Arc::new(instance.clone());
            instance
                .runtime
                .tools
                .extend(build_ability_tools(&active_abilities, base_instance)?);
        }

        debug!(
            agent = instance.name(),
            domain = domain_name,
            session_id = %instance.prompt.context.active_domain.as_ref().unwrap().session_id,
            "Domain expansion started"
        );

        Ok(Self {
            instance: Arc::new(instance),
            memory: self.memory.clone(),
            memory_scope: self.memory_scope.clone(),
            manifest: self.manifest.clone(),
        })
    }

    /// Send a chat message and stream events as the agent works.
    pub async fn chat_stream(&self, message: &str) -> Result<ExecutionHandle> {
        self.chat_with_history_stream(message, Vec::new()).await
    }

    /// Send a chat message with prior conversation history and stream events.
    pub async fn chat_with_history_stream(
        &self,
        message: &str,
        history: Vec<ChatMessage>,
    ) -> Result<ExecutionHandle> {
        self.run_stream(AgentRun::chat(ChatInput {
            message: message.to_string(),
            history,
            project: None,
        }))
        .await
    }

    /// Execute a task and stream events as the agent works.
    pub async fn task_stream(&self, task: TaskInput) -> Result<ExecutionHandle> {
        self.run_stream(AgentRun::task(task)).await
    }

    /// Execute a composed agent run and stream events as the agent works.
    pub async fn run_stream(&self, run: AgentRun) -> Result<ExecutionHandle> {
        self.execute_stream(run).await
    }

    /// Send a chat message and wait for the final output.
    pub async fn chat(&self, message: &str) -> Result<TurnOutput> {
        self.chat_stream(message).await?.output().await
    }

    /// Send a chat message with prior conversation history and wait for the final output.
    pub async fn chat_with_history(
        &self,
        message: &str,
        history: Vec<ChatMessage>,
    ) -> Result<TurnOutput> {
        self.chat_with_history_stream(message, history)
            .await?
            .output()
            .await
    }

    /// Execute a task and wait for the final output.
    pub async fn task(&self, task: TaskInput) -> Result<TurnOutput> {
        self.task_stream(task).await?.output().await
    }

    /// Execute a composed agent run and wait for the final output.
    pub async fn run(&self, run: AgentRun) -> Result<TurnOutput> {
        self.run_stream(run).await?.output().await
    }

    // -- Internal --

    async fn execute_stream(&self, run: AgentRun) -> Result<ExecutionHandle> {
        self.execute_stream_with_sub_agent(run, None).await
    }

    pub(crate) async fn task_stream_as_sub_agent(
        &self,
        task: TaskInput,
        child_handle: ChildRuntimeHandle<P>,
    ) -> Result<ExecutionHandle> {
        self.execute_stream_with_sub_agent(AgentRun::task(task), Some(child_handle))
            .await
    }

    async fn execute_stream_with_sub_agent(
        &self,
        run: AgentRun,
        child_handle: Option<ChildRuntimeHandle<P>>,
    ) -> Result<ExecutionHandle> {
        let memory_vars = if let (Some(mem), Some(scope)) = (&self.memory, &self.memory_scope)
            && self.instance.prompt.memory_vars.is_empty()
        {
            Some(memory::build_memory_vars(mem.as_ref(), scope).await?)
        } else {
            None
        };
        let artifact_vars = if let (Some(mem), Some(scope)) = (&self.memory, &self.memory_scope)
            && self.instance.prompt.artifact_vars.is_empty()
        {
            Some(memory::build_artifact_vars(mem.as_ref(), scope).await?)
        } else {
            None
        };

        // 5. Spawn the turn loop.
        let (events_tx, events_rx) = mpsc::unbounded_channel::<TurnEvent>();
        let pause_token = types::PauseToken::new();
        let loop_pause = pause_token.clone();

        let mut inst = (*self.instance).clone();
        if let Some(child_handle) = child_handle {
            inst.runtime.tools.extend(child_tools(child_handle));
        } else if inst.runtime.execution_mode == AgentExecutionMode::Parent
            && let Some(provider) = inst.runtime.provider_runtime.clone()
            && inst.runtime.config.max_delegation_depth > 0
        {
            let inherited_host_tools = inst
                .runtime
                .tools
                .iter()
                .filter(|tool| !PARENT_TOOL_NAMES.contains(&tool.name()))
                .cloned()
                .collect();
            let runtime = SubAgentRuntime::new(
                provider,
                inst.agent_id(),
                inst.model_manifest.clone(),
                inherited_host_tools,
                SubAgentLimits {
                    max_depth: inst.runtime.config.max_delegation_depth,
                },
                inst.runtime.sub_agent_ctx.clone(),
                Some(events_tx.clone()),
            );
            inst.runtime.tools.extend(parent_tools(runtime.handle()));
        }
        let inst = Arc::new(inst);

        let task_label = match &run.kind {
            AgentRunKind::Chat { .. } => "chat",
            AgentRunKind::FollowUp { .. } => "follow_up",
            AgentRunKind::Task(_) => "task",
            AgentRunKind::Heartbeat { .. } => "heartbeat",
            AgentRunKind::Gate { .. } => "gate",
        };
        let domain_label = inst
            .prompt
            .context
            .active_domain
            .as_ref()
            .map(|d| d.domain_name.as_str());

        info!(
            agent = inst.name(),
            model = inst.model.model_name.as_str(),
            task_type = task_label,
            domain = ?domain_label,
            tool_count = inst.runtime.tools.len(),
            "Executing agent"
        );

        // 3. Build prompts.
        let prompts =
            inst.build_prompts_with_vars(&run, memory_vars.as_ref(), artifact_vars.as_ref());

        let system_prompt = prompts.system;
        let developer_prompt = prompts.developer;
        let templated_user_message = prompts.user_message;
        let user_message = if !templated_user_message.is_empty() {
            templated_user_message
        } else {
            raw_user_message(&run)
        };
        trace!(
            agent = inst.name(),
            "\nRendered prompts for {}\n\n=== System Prompt ===\n{}\n\n=== Developer Prompt ===\n{}\n\n=== User Message ===\n{}",
            inst.name(),
            system_prompt,
            developer_prompt,
            user_message,
        );

        // 4. Build initial messages.
        let supports_developer_role = inst
            .model
            .model_provider
            .supports_developer_role(&inst.model.model_name);
        let mut messages: Vec<ChatMessage> =
            build_instruction_messages(&system_prompt, &developer_prompt, supports_developer_role);

        match &run.kind {
            AgentRunKind::Chat(chat) => messages.extend(chat.history.iter().cloned()),
            AgentRunKind::FollowUp(follow_up) => messages.extend(follow_up.history.iter().cloned()),
            _ => {}
        }

        messages.push(ChatMessage::user(&user_message));

        let task_id = match &run.kind {
            AgentRunKind::Task(task) => Some(task.task_id),
            AgentRunKind::Gate(crate::input::GateInput {
                task: Some(task), ..
            }) => Some(task.task_id),
            _ => None,
        };

        let join = tokio::spawn(async move {
            let mut output =
                turn_loop::run(&inst, messages, Some(events_tx), Some(loop_pause)).await?;
            output.task_id = task_id;
            Ok(output)
        });

        Ok(ExecutionHandle {
            events_rx,
            join: Some(join),
            pause_token,
        })
    }
}

fn raw_user_message(run: &AgentRun) -> String {
    match &run.kind {
        AgentRunKind::Chat(chat) => chat.message.clone(),
        AgentRunKind::FollowUp(follow_up) => follow_up.message.clone(),
        AgentRunKind::Task(task) => {
            if task.description.is_empty() {
                task.title.clone()
            } else {
                task.description.clone()
            }
        }
        AgentRunKind::Heartbeat(_) => String::new(),
        AgentRunKind::Gate(gate) => gate.previous_result.output.clone(),
    }
}

fn domain_is_assigned(assigned_domains: &[Slug], domain: &DomainManifest) -> bool {
    let manifest_slug = domain.slug();
    let name_slug = Slug::derive(&domain.name);
    assigned_domains
        .iter()
        .any(|assigned| assigned == &manifest_slug || assigned == &name_slug)
}

fn resolve_active_abilities<P: ProviderRuntime>(
    provider: &P,
    agent_abilities: &[String],
    active_domain: Option<&DomainManifest>,
) -> Vec<AbilityManifest> {
    let mut abilities = agent_abilities.to_vec();

    if let Some(domain) = active_domain {
        for ability_name in &domain.abilities {
            if !abilities.contains(ability_name) {
                abilities.push(ability_name.clone());
            }
        }
    }

    abilities
        .into_iter()
        .filter_map(|ability_name| provider.find_ability(&ability_name).cloned())
        .collect()
}

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

    #[test]
    fn instruction_messages_use_developer_when_supported() {
        let messages = build_instruction_messages("root", "app rules", true);
        assert_eq!(messages.len(), 1);
        assert_eq!(messages[0].role, "developer");
        assert_eq!(messages[0].content, "root\n\napp rules");
    }

    #[test]
    fn instruction_messages_fallback_to_system_when_developer_unsupported() {
        let messages = build_instruction_messages("root", "app rules", false);
        assert_eq!(messages.len(), 1);
        assert_eq!(messages[0].role, "system");
        assert_eq!(messages[0].content, "root\n\napp rules");
    }
}
