//! The orchestrator: the assistant threads and the channels that join them.
//!
//! ```text
//!   microphone ──samples──> segmenter ──utterance──> ASR ──text──┐
//!   (cpal, real time)       (VAD, turns)             (Canary)    │
//!                                                                 v
//!   speaker <──samples── synthesis <──sentences── conversation <─┘
//!   (cpal, ring)          (qwentts)                (llama.cpp + tools)
//! ```
//!
//! Each box is a thread and each arrow a channel. Two consequences carry the
//! whole design:
//!
//! * **Nothing blocks the stage ahead.** The microphone never waits for the
//!   ASR, and the model never waits for the synthesizer; whoever has spare
//!   capacity drops work instead of accumulating delay.
//! * **Answers go out per sentence, not per answer.** The conversation hands
//!   each sentence to the synthesizer as soon as it is closed, so the assistant
//!   starts speaking while the model is still writing. That is what separates
//!   half a second of latency from five.

use std::sync::Arc;
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use crossbeam_channel::{bounded, unbounded, Receiver, Select, Sender};

use asist_asr::{AsrJob, AsrResult, Recognizer};
use asist_audio::{
    CaptureBlock, Gate, InputFormat, PlaybackHandle, Segmenter, VoiceEvent, ASR_SAMPLE_RATE,
};
use asist_core::config::Config;
use asist_core::error::Result;
use asist_core::event::{Event, InterruptReason, TurnId};
use asist_core::telemetry::{Metrics, Stage, TurnTimer};
use asist_core::text::SentenceSplitter;
use asist_core::tools::ToolRegistry;
use asist_llm::chat::Conversation;
use asist_llm::{Delta, LlmClient, Message};
use asist_tts::TtsClient;

use crate::session::Session;

/// A sentence waiting to be synthesized.
#[derive(Debug)]
pub struct SpeakJob {
    pub turn: TurnId,
    pub index: usize,
    pub text: String,
    /// Origin from which the perceived latency of the turn is measured.
    pub turn_started: Instant,
}

/// The running threads, so they can be joined on shutdown.
pub struct Pipeline {
    handles: Vec<JoinHandle<()>>,
    pub events: Receiver<Event>,
    pub metrics: Arc<Metrics>,
    session: Session,
    capture_tx: Option<Sender<CaptureBlock>>,
}

impl Pipeline {
    /// Builds the whole pipeline. `capture_rx` comes from the microphone.
    #[allow(clippy::too_many_arguments)]
    pub fn spawn(
        config: &Config,
        session: Session,
        recognizer: Recognizer,
        llm: Arc<LlmClient>,
        tts: TtsClient,
        tools: ToolRegistry,
        playback: PlaybackHandle,
        capture_rx: Receiver<CaptureBlock>,
        capture_tx: Sender<CaptureBlock>,
        input_format: InputFormat,
    ) -> Result<Self> {
        let (event_tx, event_rx) = unbounded::<Event>();
        let (asr_tx, asr_rx) = unbounded::<AsrJob>();
        let (result_tx, result_rx) = unbounded::<AsrResult>();
        // Bounded on purpose: if synthesis gets stuck, the conversation thread
        // must notice and slow down instead of piling up sentences that will
        // arrive late to a turn that may already have been interrupted.
        let (speak_tx, speak_rx) = bounded::<SpeakJob>(8);

        let metrics = Arc::new(Metrics::default());
        let mut handles = Vec::new();

        handles.push(spawn_named("segmenter", {
            let config = config.clone();
            let session = session.clone();
            let events = event_tx.clone();
            let playback = playback.clone();
            move || {
                run_segmenter(
                    &config,
                    session,
                    capture_rx,
                    asr_tx,
                    events,
                    playback,
                    input_format,
                )
            }
        }));

        handles.push(spawn_named("asr", move || {
            recognizer.run(asr_rx, result_tx)
        }));

        handles.push(spawn_named("conversation", {
            let config = config.clone();
            let session = session.clone();
            let events = event_tx.clone();
            let metrics = Arc::clone(&metrics);
            move || {
                run_brain(
                    &config, session, llm, tools, result_rx, speak_tx, events, metrics,
                )
            }
        }));

        handles.push(spawn_named("synthesis", {
            let session = session.clone();
            let events = event_tx;
            move || run_speech(session, tts, speak_rx, playback, events)
        }));

        Ok(Self {
            handles,
            events: event_rx,
            metrics,
            session,
            capture_tx: Some(capture_tx),
        })
    }

    /// Shuts down in order: what is playing is cut, the microphone sender is
    /// dropped, and with that the chain of threads takes itself apart.
    pub fn shutdown(mut self) {
        self.session.request_stop();
        self.capture_tx.take();
        for handle in self.handles.drain(..) {
            let _ = handle.join();
        }
    }
}

fn spawn_named(name: &str, body: impl FnOnce() + Send + 'static) -> JoinHandle<()> {
    thread::Builder::new()
        .name(name.to_string())
        .spawn(body)
        .expect("could not create a pipeline thread")
}

// ---------------------------------------------------------------------------
// Segmenter: microphone -> utterances
// ---------------------------------------------------------------------------

fn run_segmenter(
    config: &Config,
    session: Session,
    capture: Receiver<CaptureBlock>,
    asr: Sender<AsrJob>,
    events: Sender<Event>,
    playback: PlaybackHandle,
    input_format: InputFormat,
) {
    let mut segmenter = Segmenter::new(&config.vad);
    let mut turn = session.current();

    while let Ok(block) = capture.recv() {
        if session.is_stopping() {
            break;
        }
        // The microphone is closed while the speaker is playing. With barge-in
        // on it is not closed; only the volume bar goes up.
        //
        // The second condition looks at the queue and not at a flag on purpose:
        // the queue cannot fall out of sync, and a flag someone forgets to lower
        // leaves the microphone closed for the rest of the session.
        segmenter.set_gate(if session.is_speaking() || playback.queued_secs() > 0.0 {
            Gate::Speaking
        } else {
            Gate::Open
        });

        // The device is opened at 16 kHz mono whenever possible, and then this
        // does nothing; when it cannot be, this is where it gets converted.
        let mono = asist_audio::to_mono_at(&block.samples, input_format, ASR_SAMPLE_RATE);

        for event in segmenter.push(&mono) {
            match event {
                VoiceEvent::BargeIn => {
                    // Cut before opening the new turn: that way whatever was left
                    // in the ring does not leak over it.
                    playback.stop();
                    session.interrupt();
                    let _ = events.send(Event::Interrupted {
                        turn: session.current(),
                        reason: InterruptReason::UserSpoke,
                    });
                }
                VoiceEvent::Started => {
                    turn = session.begin_turn();
                    let _ = events.send(Event::SpeechStarted { turn, at: block.at });
                }
                VoiceEvent::Audio(samples) => {
                    if config.asr.partials
                        && asr
                            .send(AsrJob::Window {
                                turn,
                                samples,
                                at: block.at,
                            })
                            .is_err()
                    {
                        return;
                    }
                }
                VoiceEvent::Ended(utterance) => {
                    let at = Instant::now();
                    if asr
                        .send(AsrJob::Utterance {
                            turn,
                            samples: utterance.samples,
                            at,
                        })
                        .is_err()
                    {
                        return;
                    }
                }
                VoiceEvent::Discarded => {
                    if asr.send(AsrJob::Reset { turn }).is_err() {
                        return;
                    }
                }
            }
        }
    }
    // Whatever was left half-done at shutdown still deserves a transcription.
    if let Some(utterance) = segmenter.flush() {
        let _ = asr.send(AsrJob::Utterance {
            turn,
            samples: utterance.samples,
            at: Instant::now(),
        });
    }
}

// ---------------------------------------------------------------------------
// Conversation: transcription -> sentences
// ---------------------------------------------------------------------------

#[allow(clippy::too_many_arguments)]
fn run_brain(
    config: &Config,
    session: Session,
    llm: Arc<LlmClient>,
    tools: ToolRegistry,
    results: Receiver<AsrResult>,
    speak: Sender<SpeakJob>,
    events: Sender<Event>,
    metrics: Arc<Metrics>,
) {
    let prompts = Prompts::new(config, &tools);
    let mut chat = Conversation::new(prompts.deciding.clone(), config.general.history_turns);

    while let Ok(result) = results.recv() {
        if session.is_stopping() {
            break;
        }
        match result {
            AsrResult::Partial {
                turn,
                committed,
                volatile,
                ..
            } => {
                let _ = events.send(Event::Partial {
                    turn,
                    committed,
                    volatile,
                });
            }
            AsrResult::Empty { turn } => {
                let _ = events.send(Event::Discarded { turn });
            }
            AsrResult::Error { turn, message } => {
                let _ = events.send(Event::Warning { turn, message });
            }
            AsrResult::Final {
                turn,
                text,
                audio_secs,
                decode,
                spoken_at,
            } => {
                // A transcription from an outdated turn arrives too late: the
                // user has already said something else.
                if !session.is_current(turn) {
                    tracing::debug!(target: "brain", turn = turn.0, "stale transcription");
                    continue;
                }
                let mut timer = TurnTimer::new(turn);
                timer.input_secs = audio_secs;
                timer.record(Stage::Asr, decode);

                let _ = events.send(Event::Transcript {
                    turn,
                    text: text.clone(),
                    audio_secs,
                    decode,
                });

                chat.push(Message::user(text));
                answer(
                    config, &session, &llm, &tools, &prompts, &mut chat, &speak, &events,
                    &mut timer, spoken_at,
                );

                if config.general.report_latency {
                    tracing::info!(target: "latency", "\n{}", timer.report());
                }
                metrics.record_turn(&timer);
            }
        }
    }
}

/// The two system prompts a turn alternates between.
///
/// Having two is not a design whim but a measured result: with Qwen3.5-2B,
/// any style instruction next to the tool guide makes the model stop calling
/// tools and make the data up (8/8 hits with the guide alone, 0/8 with the
/// voice-assistant persona added). So the pass where the model *decides*
/// whether to act carries the guide alone, and the voice prompt is kept for
/// writing what will be spoken.
struct Prompts {
    /// For the pass where the model decides whether to call a tool.
    deciding: String,
    /// For writing the spoken answer, with the results already at hand.
    ///
    /// It carries the instruction to use what the tool returned. Without it the
    /// model just announces what it did («he tomado una foto, ahora puedo
    /// responderte sobre el objeto o color») and leaves out the data it had in
    /// front of it. Style can be added safely here: the call already happened,
    /// so there is nothing left to break.
    speaking: String,
    /// `false` when there are no tools: then both are the same.
    two_phase: bool,
}

impl Prompts {
    fn new(config: &Config, tools: &ToolRegistry) -> Self {
        let speaking = config.general.system_prompt.trim().to_string();
        let two_phase = config.tools.enabled && config.tools.dedicated_prompt && !tools.is_empty();
        Self {
            deciding: if two_phase {
                config.general.tools_prompt.trim().to_string()
            } else {
                speaking.clone()
            },
            speaking: if two_phase {
                format!("{speaking}\n\n{}", config.general.tool_result_prompt.trim())
            } else {
                speaking.clone()
            },
            two_phase,
        }
    }
}

/// Generates the answer of a turn, resolving tools if the model asks for
/// them, and releases sentences to the synthesizer as they close.
#[allow(clippy::too_many_arguments)]
fn answer(
    config: &Config,
    session: &Session,
    llm: &LlmClient,
    tools: &ToolRegistry,
    prompts: &Prompts,
    chat: &mut Conversation,
    speak: &Sender<SpeakJob>,
    events: &Sender<Event>,
    timer: &mut TurnTimer,
    turn_started: Instant,
) {
    let turn = timer.turn;
    let tools = (!tools.is_empty() && config.tools.enabled).then_some(tools);
    let mut splitter = SentenceSplitter::new();
    let mut spoken = String::new();
    // Own index, not the chunker's, because tool acknowledgements slip into
    // the sentence stream. Index 0 marks the first sound of the turn, which is
    // where perceived latency is measured from.
    let mut emitted = 0usize;

    // Sends a sentence to the synthesizer. Returns `false` if the channel closed.
    let say = |index: &mut usize, text: String| -> bool {
        let _ = events.send(Event::Sentence {
            turn,
            index: *index,
            text: text.clone(),
        });
        let job = SpeakJob {
            turn,
            index: *index,
            text,
            turn_started,
        };
        *index += 1;
        speak.send(job).is_ok()
    };

    for round in 0..=config.llm.max_tool_rounds {
        let llm_started = Instant::now();
        let mut first_token = None;

        let outcome = llm.stream(chat, tools, session.llm_cancel(), |delta| {
            if !session.is_current(turn) {
                return false;
            }
            match delta {
                Delta::Text(text) => {
                    if first_token.is_none() {
                        first_token = Some(llm_started.elapsed());
                        let _ = events.send(Event::ReplyStarted {
                            turn,
                            ttft: llm_started.elapsed(),
                        });
                    }
                    let _ = events.send(Event::ReplyDelta {
                        turn,
                        text: text.clone(),
                    });
                    spoken.push_str(text);
                    // This is the overlap: every closed sentence goes to the
                    // synthesizer without waiting for the rest of the answer.
                    for sentence in splitter.push(text) {
                        if !say(&mut emitted, sentence) {
                            return false;
                        }
                    }
                    true
                }
                Delta::ToolCalls(_) => true,
            }
        });

        if let Some(ttft) = first_token {
            timer.record(Stage::LlmFirstToken, ttft);
            timer.record(Stage::LlmRest, llm_started.elapsed().saturating_sub(ttft));
        }

        let outcome = match outcome {
            Ok(outcome) => outcome,
            Err(err) => {
                let _ = events.send(Event::Failed {
                    turn,
                    message: err.to_string(),
                });
                return;
            }
        };

        if !session.is_current(turn) {
            return;
        }

        if outcome.tool_calls.is_empty() {
            if let Some(rest) = splitter.flush() {
                say(&mut emitted, rest);
            }
            timer.sentences = emitted;
            chat.push(Message::assistant(outcome.text.clone()));
            // The next turn starts again by deciding.
            if prompts.two_phase {
                chat.set_system(&prompts.deciding);
            }
            let _ = events.send(Event::ReplyDone {
                turn,
                text: outcome.text,
            });
            return;
        }

        // The model asked for tools: they run, the result goes back to it and
        // it is asked again.
        if round == config.llm.max_tool_rounds {
            let _ = events.send(Event::Warning {
                turn,
                message: format!(
                    "the model kept asking for tools after {} rounds; stopping",
                    config.llm.max_tool_rounds
                ),
            });
            return;
        }

        let tools = tools.expect("there cannot be calls without registered tools");
        chat.push(Message::tool_request(
            outcome.text.clone(),
            outcome.tool_calls.clone(),
        ));
        let tools_started = Instant::now();
        for call in &outcome.tool_calls {
            let _ = events.send(Event::ToolRequested {
                turn,
                name: call.name.clone(),
                arguments: call.arguments.clone(),
            });
            // Said before running, not after: the point is to cover the
            // wait, and a search with the two model passes behind it is
            // six seconds that would otherwise pass in total silence.
            if config.tools.spoken_ack {
                if let Some(ack) = tools.get(&call.name).and_then(|t| t.acknowledgement()) {
                    say(&mut emitted, ack.to_string());
                }
            }
            let result = tools.dispatch(call);
            let _ = events.send(Event::ToolFinished {
                turn,
                name: result.name.clone(),
                ok: result.ok,
                output: result.output.clone(),
                took: result.took,
            });
            chat.push(Message::tool_result(&result));
            timer.tool_calls += 1;
        }
        timer.record(Stage::Tools, tools_started.elapsed());

        // With the result already in the conversation, the next round only
        // has to write: this is when the voice prompt comes back, which in
        // the previous pass would have prevented the call.
        if prompts.two_phase {
            chat.set_system(&prompts.speaking);
        }
    }
}

// ---------------------------------------------------------------------------
// Synthesis: sentences -> speaker
// ---------------------------------------------------------------------------

fn run_speech(
    session: Session,
    tts: TtsClient,
    jobs: Receiver<SpeakJob>,
    playback: PlaybackHandle,
    events: Sender<Event>,
) {
    // A `Select` instead of a plain `recv()` so it can wake up periodically
    // and lower the «speaking» flag when the ring empties: otherwise the
    // microphone would stay closed after the last sentence.
    let mut select = Select::new();
    let job_index = select.recv(&jobs);
    let mut speaking_turn: Option<TurnId> = None;

    loop {
        let job = match select.select_timeout(Duration::from_millis(100)) {
            Ok(op) if op.index() == job_index => match op.recv(&jobs) {
                Ok(job) => Some(job),
                Err(_) => break,
            },
            Ok(_) => None,
            Err(_) => None,
        };

        let Some(job) = job else {
            // No work: if there is no audio left, the turn has finished playing.
            if let Some(turn) = speaking_turn {
                if playback.queued_secs() <= 0.0 {
                    session.set_speaking(false);
                    playback.mark_idle();
                    speaking_turn = None;
                    let _ = events.send(Event::AudioFinished { turn });
                }
            }
            if session.is_stopping() {
                break;
            }
            continue;
        };

        // A sentence from an outdated turn must never be heard.
        if !session.is_current(job.turn) {
            tracing::debug!(target: "tts", turn = job.turn.0, "stale sentence, dropped");
            continue;
        }

        session.set_speaking(true);
        speaking_turn = Some(job.turn);
        if job.index == 0 {
            playback.reset_played();
        }

        let turn = job.turn;
        let first_of_turn = job.index == 0;
        let mut announced = false;
        let result = tts.speak(&job.text, session.tts_cancel(), |samples| {
            if !session.is_current(turn) {
                return false;
            }
            playback.push_tts(samples);
            if first_of_turn && !announced && playback.has_played() {
                announced = true;
                let _ = events.send(Event::AudioStarted {
                    turn,
                    latency: job.turn_started.elapsed(),
                });
            }
            true
        });

        match result {
            Ok(outcome) => {
                if outcome.cancelled {
                    let _ = events.send(Event::Interrupted {
                        turn,
                        reason: InterruptReason::UserSpoke,
                    });
                    continue;
                }
                // The first-audio mark may not have been set if the ring had
                // not played anything yet when the block arrived.
                if first_of_turn && !announced {
                    let _ = events.send(Event::AudioStarted {
                        turn,
                        latency: job.turn_started.elapsed(),
                    });
                }
                if outcome.rtf() > 1.0 {
                    tracing::warn!(
                        target: "tts",
                        rtf = outcome.rtf(),
                        "synthesis is running behind real time; the voice will break up"
                    );
                }
            }
            Err(err) => {
                let _ = events.send(Event::Failed {
                    turn,
                    message: err.to_string(),
                });
            }
        }
    }

    playback.stop();
    session.set_speaking(false);
}