//! Streaming execution for the Agent.
//!
//! This module provides streaming variants of agent execution that yield
//! events as they occur, enabling real-time feedback during agent runs.
//!
//! # Parallel Tool Execution
//!
//! When multiple tool calls are returned by the model, they are executed
//! in parallel (respecting `max_parallel_tool_calls` config). Events are
//! yielded in the following order:
//!
//! 1. `ToolCallStart` for each tool (in order)
//! 2. Tools execute in parallel
//! 3. `ToolCallComplete` for each tool (as they complete)

use std::sync::atomic::Ordering;

use async_stream::stream;
use futures::{Stream, StreamExt};
use serde_json::Value;
use tracing::{debug, info, warn};

use crate::{
    error::AgentError,
    memory::{ActionStep, FinalAnswerStep, Timing, ToolCall},
    message::ChatMessageStreamDelta,
    providers::common::GenerateOptions,
    tools::FinalAnswerArgs,
};

use super::{
    Agent,
    events::{StepResult, StreamEvent, StreamItem},
    tool_processor::ToolProcessor,
};

impl Agent {
    /// Stream execution events with full options.
    ///
    /// This is the core streaming implementation that yields events at both
    /// step-level and token-level granularity. TextDelta events are yielded
    /// in real-time as tokens are generated by the model.
    #[expect(
        tail_expr_drop_order,
        reason = "stream yields control flow intentionally"
    )]
    pub(crate) fn stream_execution(&mut self) -> impl Stream<Item = StreamItem> + '_ {
        info!("Starting streaming agent run");

        stream! {
            let mut final_answer: Option<Value> = None;

            while self.step_number < self.config.max_steps {
                // Check for interruption
                if self.interrupt_flag.load(Ordering::SeqCst) {
                    yield Err(AgentError::Interrupted);
                    break;
                }

                self.step_number += 1;
                let mut step = ActionStep {
                    step_number: self.step_number,
                    timing: Timing::start_now(),
                    ..Default::default()
                };

                // Prepare messages and options
                let messages = self.memory.to_messages(false);
                step.model_input_messages = Some(messages.clone());
                let options = GenerateOptions::new().with_tools(self.tools.definitions());
                debug!(step = step.step_number, "Generating model response");

                // Stream model response with real-time token output
                let mut deltas: Vec<ChatMessageStreamDelta> = Vec::new();
                let model_stream_result = self.model.generate_stream(messages, options).await;

                let message = match model_stream_result {
                    Ok(mut model_stream) => {
                        while let Some(result) = model_stream.next().await {
                            match result {
                                Ok(delta) => {
                                    // Yield text delta immediately for real-time output
                                    if let Some(content) = &delta.content
                                        && !content.is_empty()
                                    {
                                        yield Ok(StreamEvent::TextDelta(content.clone()));
                                    }
                                    if let Some(usage) = &delta.token_usage {
                                        step.token_usage = Some(*usage);
                                    }
                                    deltas.push(delta);
                                }
                                Err(e) => {
                                    step.error = Some(e.to_string());
                                    yield Err(AgentError::internal(e.to_string()));
                                    break;
                                }
                            }
                        }
                        let msg = crate::message::aggregate_stream_deltas(&deltas);
                        step.model_output_message = Some(msg.clone());
                        step.model_output = msg.text_content();
                        Some(msg)
                    }
                    Err(e) => {
                        step.error = Some(e.to_string());
                        yield Ok(StreamEvent::Error(e.to_string()));
                        None
                    }
                };

                // Process tool calls if we got a message
                let step_outcome = if let Some(ref msg) = message {
                    // Extract tool calls from message or parse from text
                    let tool_calls = msg.tool_calls.clone().or_else(|| {
                        step.model_output.as_ref().and_then(|text| {
                            ToolProcessor::parse_text_tool_call(text).map(|tc| vec![tc])
                        })
                    });

                    if let Some(calls) = tool_calls {
                        // Separate final_answer from regular tool calls
                        let mut got_final = None;
                        let mut regular_calls: Vec<(&str, String, Value)> = Vec::new();

                        // First pass: record all tool calls and yield start events
                        for tc in &calls {
                            let tool_name = tc.name();
                            let tool_id = tc.id.clone();

                            // Record in step memory
                            step.tool_calls.get_or_insert_with(Vec::new).push(
                                ToolCall::new(&tool_id, tool_name, tc.arguments().clone())
                            );

                            // Yield start event
                            yield Ok(StreamEvent::ToolCallStart {
                                id: tool_id.clone(),
                                name: tool_name.to_string(),
                            });

                            if tool_name == "final_answer" {
                                let answer = tc.parse_arguments::<FinalAnswerArgs>()
                                    .map_or_else(|_| tc.arguments().clone(), |args| args.answer);
                                got_final = Some(answer);
                                step.is_final_answer = true;
                                yield Ok(StreamEvent::ToolCallComplete {
                                    id: tool_id,
                                    name: tool_name.to_string(),
                                    result: Ok("Final answer recorded".to_string()),
                                });
                            } else {
                                // Queue for parallel execution
                                regular_calls.push((tool_name, tool_id, tc.arguments().clone()));
                            }
                        }

                        // Execute regular tools in parallel
                        let observations = if regular_calls.is_empty() {
                            Vec::new()
                        } else {
                            debug!(
                                count = regular_calls.len(),
                                max_concurrent = ?self.config.max_parallel_tool_calls,
                                "Executing tool calls in parallel"
                            );

                            let results = self.tools.call_parallel(
                                regular_calls,
                                self.config.max_parallel_tool_calls,
                            ).await;

                            let mut obs = Vec::with_capacity(results.len());
                            for result in results {
                                let observation = result.to_observation();

                                // Yield completion event - convert Result<Value, ToolError> to Result<String, String>
                                let result_str: Result<String, String> = match &result.result {
                                    Ok(v) => Ok(v.to_string()),
                                    Err(e) => Err(e.to_string()),
                                };
                                yield Ok(StreamEvent::ToolCallComplete {
                                    id: result.id.clone(),
                                    name: result.name.clone(),
                                    result: result_str,
                                });

                                if result.is_err() {
                                    step.error = Some(observation.clone());
                                }

                                obs.push(observation);
                            }
                            obs
                        };

                        if !observations.is_empty() {
                            step.observations = Some(observations.join("\n"));
                        }

                        match got_final {
                            Some(answer) => {
                                step.action_output = Some(answer.clone());
                                Ok(StepResult::FinalAnswer(answer))
                            }
                            None => Ok(StepResult::Continue),
                        }
                    } else {
                        Ok(StepResult::Continue)
                    }
                } else {
                    Err(AgentError::internal("No model response"))
                };

                // Finalize step timing and telemetry
                step.timing.complete();
                self.record_telemetry(&step);

                // Invoke callbacks
                let ctx = self.create_callback_context();
                self.callbacks.callback(&step, &ctx);

                // Yield step complete event
                let step_clone = step.clone();
                self.memory.add_step(step);
                yield Ok(StreamEvent::StepComplete {
                    step: self.step_number,
                    action_step: Box::new(step_clone),
                });

                // Handle step result
                match step_outcome {
                    Ok(StepResult::FinalAnswer(answer)) => {
                        if let Err(e) = self.validate_answer(&answer) {
                            warn!(error = %e, "Final answer check failed");
                            yield Ok(StreamEvent::Error(format!("Final answer check failed: {e}")));
                            continue;
                        }

                        let final_step = FinalAnswerStep { output: answer.clone() };
                        self.callbacks.callback(&final_step, &ctx);
                        self.memory.add_step(FinalAnswerStep { output: answer.clone() });

                        final_answer = Some(answer.clone());
                        yield Ok(StreamEvent::FinalAnswer { answer });
                        break;
                    }
                    Ok(StepResult::Continue) => {}
                    Err(e) => {
                        warn!(step = self.step_number, error = %e, "Step failed");
                        yield Ok(StreamEvent::Error(e.to_string()));
                    }
                }
            }

            // Handle max steps reached
            if final_answer.is_none() && self.step_number >= self.config.max_steps {
                let error_msg = format!("Maximum steps ({}) reached", self.config.max_steps);
                self.memory.add_step(FinalAnswerStep {
                    output: Value::String(error_msg.clone()),
                });
                yield Err(AgentError::max_steps(self.step_number, self.config.max_steps));
            }
        }
    }
}
