//! Tests against the real servers.
//!
//! They skip themselves when nothing is listening, so `cargo test` stays
//! useful without 5 GB of models loaded. With the servers up they check what
//! unit tests cannot: that the protocol written here is the one the servers
//! speak, and that latencies are still where they were measured.
//!
//! To run them:  scripts/servers.sh start && cargo test -- --ignored

use std::path::PathBuf;
use std::sync::{Mutex, MutexGuard, OnceLock};
use std::time::{Duration, Instant};

use asist_core::config::Config;
use asist_core::http::Cancel;
use asist_core::tools::{Tool, ToolRegistry};
use asist_llm::chat::Conversation;
use asist_llm::{Delta, LlmClient, Message};
use asist_tools::FrameSource;
use asist_tts::TtsClient;

/// The tests take turns on the GPU.
///
/// Both servers share a 4 GB card, and measured: with the synthesis test
/// running at the same time, the LLM first token goes from ~0.5 s to ~5 s.
/// In parallel this does not measure latency, it measures who got there
/// first. The real pipeline does not have this problem because synthesis
/// does not start until the model closes its first sentence.
fn exclusive() -> MutexGuard<'static, ()> {
    static GPU: OnceLock<Mutex<()>> = OnceLock::new();
    GPU.get_or_init(|| Mutex::new(()))
        .lock()
        .unwrap_or_else(|e| e.into_inner())
}

fn config() -> Config {
    let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../config/asistente.toml");
    Config::load(path).expect("could not load the configuration")
}

/// Returns the clients, or `None` if the servers are not up.
fn clients() -> Option<(Config, LlmClient, TtsClient)> {
    let config = config();
    let llm = LlmClient::new(config.llm_authority(), &config.llm);
    let tts = TtsClient::new(config.tts_authority(), &config.tts);
    if !llm.healthy() || !tts.healthy() {
        eprintln!("servidores no disponibles; prueba omitida");
        return None;
    }
    Some((config, llm, tts))
}

#[test]
fn llm_streams_without_thinking_aloud() {
    let _gpu = exclusive();
    let Some((config, llm, _)) = clients() else {
        return;
    };

    let mut chat = Conversation::new(&config.general.system_prompt, 4);
    chat.push(Message::user("Saluda en una frase corta."));

    let started = Instant::now();
    let mut ttft = None;
    let outcome = llm
        .stream(&chat, None, &Cancel::new(), |delta| {
            if let Delta::Text(_) = delta {
                ttft.get_or_insert_with(|| started.elapsed());
            }
            true
        })
        .expect("the LLM request failed");

    assert!(!outcome.text.trim().is_empty(), "empty answer");
    // The template with an unclosed <think> puts this at 7-9 s. If this test
    // fails, llama-server was almost certainly started without --chat-template-file.
    let ttft = ttft.expect("no fragment with text arrived");
    assert!(
        ttft < Duration::from_secs(3),
        "the first token took {ttft:?}: was llama-server started with the no-reasoning template?"
    );
    assert!(
        !outcome.text.contains("<think>"),
        "reasoning is leaking into the answer: {}",
        outcome.text
    );
}

#[test]
fn llm_can_request_a_tool() {
    let _gpu = exclusive();
    let Some((config, llm, _)) = clients() else {
        return;
    };

    let tools = ToolRegistry::from_config(&config.tools);
    let mut chat = Conversation::new(
        "Eres un asistente. Usa las herramientas disponibles cuando hagan falta.",
        4,
    );
    chat.push(Message::user(
        "¿Qué hora es exactamente? Usa la herramienta.",
    ));

    let outcome = llm
        .stream(&chat, Some(&tools), &Cancel::new(), |_| true)
        .expect("the LLM request failed");

    if outcome.tool_calls.is_empty() {
        // A 2B model does not always manage to call; what is checked is that the
        // round trip works, not that the model is smart.
        eprintln!("the model did not ask for a tool: {}", outcome.text);
        return;
    }
    let call = &outcome.tool_calls[0];
    assert_eq!(call.name, "hora_actual");
    let result = tools.dispatch(call);
    assert!(result.ok, "the tool failed: {}", result.output);
    assert!(!result.output.trim().is_empty());
}

#[test]
fn synthesis_starts_playing_before_it_finishes() {
    let _gpu = exclusive();
    let Some((config, _, tts)) = clients() else {
        return;
    };

    if let Some(reference) = &config.tts.reference {
        tts.register_voice(reference)
            .expect("the voice was not registered");
    }
    // Without warmup, the first synthesis loads the graphs and measures the
    // server startup instead of the steady state.
    tts.warmup().expect("warmup failed");

    let mut blocks = 0usize;
    let outcome = tts
        .speak(
            "Hola, esto es una prueba de latencia del sintetizador de voz.",
            &Cancel::new(),
            |samples| {
                if !samples.is_empty() {
                    blocks += 1;
                }
                true
            },
        )
        .expect("synthesis failed");

    assert!(outcome.samples > 0, "no audio received");
    assert!(
        blocks > 1,
        "the audio arrived all at once ({blocks} block): check --codec-chunk-dur, \
         whose 24 s default blocks streaming"
    );
    let ttfb = outcome.ttfb.expect("no first-block mark");
    assert!(
        ttfb < Duration::from_secs(2),
        "first audio took {ttfb:?}; 585 ms was measured with --codec-chunk-dur 1.0"
    );
    eprintln!(
        "ttfb={:?} audio={:.2}s rtf={:.2}",
        ttfb,
        outcome.audio_secs(),
        outcome.rtf()
    );
}

#[test]
fn synthesis_can_be_cut_mid_sentence() {
    let _gpu = exclusive();
    let Some((config, _, tts)) = clients() else {
        return;
    };
    if let Some(reference) = &config.tts.reference {
        let _ = tts.register_voice(reference);
    }

    // It is the operation barge-in relies on: if it could not be cut, the
    // assistant would keep talking over the user until the sentence ended.
    let cancel = Cancel::new();
    let mut received = 0usize;
    let outcome = tts
        .speak(
            "Esta es una frase larga que no debería llegar a escucharse entera \
             porque se va a cortar en cuanto empiece a sonar el primer bloque.",
            &cancel,
            |samples| {
                received += samples.len();
                // Cut at the first block.
                false
            },
        )
        .expect("synthesis failed");

    assert!(outcome.cancelled, "it should have been marked as cancelled");
    assert!(
        outcome.audio_secs() < 3.0,
        "{:.1} s of audio received: the cut had no effect",
        outcome.audio_secs()
    );
}

#[test]
fn cloned_voice_gets_registered() {
    let _gpu = exclusive();
    let Some((config, _, tts)) = clients() else {
        return;
    };
    let Some(reference) = &config.tts.reference else {
        eprintln!("no reference voice configured; test skipped");
        return;
    };
    tts.register_voice(reference)
        .expect("the voice was not registered");
    let voices = tts.voices().expect("could not list the voices");
    assert!(
        voices.contains(&reference.name),
        "«{}» no aparece entre {voices:?}",
        reference.name
    );
}

#[test]
#[ignore = "loads the ASR model (~200 MB) and decodes; slow"]
fn asr_transcribes_what_tts_synthesizes() {
    let _gpu = exclusive();
    // The full loop without a microphone: a known sentence is synthesized and
    // the recognizer is checked to get it back. It is the only test that
    // exercises ASR and TTS on the same audio.
    let Some((config, _, tts)) = clients() else {
        return;
    };
    if let Some(reference) = &config.tts.reference {
        let _ = tts.register_voice(reference);
    }

    let sentence = "hola qué tal estás hoy";
    let mut audio: Vec<f32> = Vec::new();
    tts.speak(sentence, &Cancel::new(), |samples| {
        audio.extend_from_slice(samples);
        true
    })
    .expect("synthesis failed");
    assert!(!audio.is_empty(), "no audio was generated");

    let recognizer =
        asist_asr::Recognizer::load(&config.asr).expect("could not load the ASR model");
    let text_in = recognizer
        .transcribe(&audio, asist_tts::SAMPLE_RATE)
        .expect("transcription failed")
        .to_lowercase();

    eprintln!("said: «{sentence}» / heard: «{text_in}»");
    // An exact match is not required (24 kHz resampled to 16 kHz and a 180M
    // model only go so far), just that it recognizes something.
    let hits = sentence
        .split_whitespace()
        .filter(|word| text_in.contains(word))
        .count();
    assert!(
        hits >= 2,
        "only {hits} words of «{sentence}» were recognized in «{text_in}»"
    );
}

// ---------------------------------------------------------------------------
// Web search and camera
// ---------------------------------------------------------------------------

#[test]
fn search_returns_something_speakable() {
    let config = config();
    if !config.search.enabled {
        eprintln!("search disabled; test skipped");
        return;
    }
    let Ok(key) = std::env::var(&config.search.api_key_env) else {
        eprintln!("no ${}; test skipped", config.search.api_key_env);
        return;
    };
    if key.trim().is_empty() {
        eprintln!("${} is empty; test skipped", config.search.api_key_env);
        return;
    }

    let tool = asist_tools::WebSearch::new(
        asist_tools::SearchBackend::Tavily { api_key: key },
        config.search.max_results,
        Duration::from_secs(config.search.timeout_secs),
    );
    let started = Instant::now();
    let out = tool
        .call(&serde_json::json!({ "consulta": "capital de Australia" }))
        .expect("the search failed");

    eprintln!(
        "search in {:?}: {}",
        started.elapsed(),
        &out[..out.len().min(200)]
    );
    assert!(!out.trim().is_empty());
    assert!(
        out.to_lowercase().contains("canberra"),
        "expected the answer in the summary, got: {out}"
    );
    // What it returns is read aloud; a wall of thousands of characters drowns
    // the model that has to summarize it.
    assert!(
        out.len() < 6000,
        "the result is too long: {} bytes",
        out.len()
    );
}

#[test]
fn made_up_query_does_not_break_the_turn() {
    let config = config();
    let Ok(key) = std::env::var(&config.search.api_key_env) else {
        return;
    };
    if key.trim().is_empty() {
        return;
    }
    // Through the registry, which is how it really arrives: a failure must come
    // back as text for the model, never as an error that cuts the answer.
    let mut registry = ToolRegistry::new();
    registry.register(std::sync::Arc::new(asist_tools::WebSearch::new(
        asist_tools::SearchBackend::Tavily {
            api_key: "clave-invalida".into(),
        },
        3,
        Duration::from_secs(10),
    )));
    let outcome = registry.dispatch(&asist_core::tools::ToolCall {
        id: "1".into(),
        name: "buscar_en_internet".into(),
        arguments: r#"{"consulta":"algo"}"#.into(),
    });
    assert!(!outcome.ok);
    assert!(
        !outcome.output.contains("clave-invalida"),
        "the error message must not contain the key: {}",
        outcome.output
    );
}

fn camera_source(config: &Config) -> Option<asist_tools::Camera> {
    if !config.camera.enabled || !config.camera.device.exists() {
        eprintln!("no camera; test skipped");
        return None;
    }
    Some(asist_tools::Camera::new(asist_tools::CameraConfig {
        device: config.camera.device.clone(),
        width: config.camera.width,
        height: config.camera.height,
        warmup_frames: config.camera.warmup_frames,
        timeout: Duration::from_secs(config.camera.timeout_secs),
        save_dir: None,
    }))
}

fn screen_source(config: &Config) -> Option<asist_tools::Screen> {
    let screen = asist_tools::Screen::new(asist_tools::ScreenConfig {
        command: config.screen.command.clone(),
        width: config.screen.width,
        output: config.screen.output.clone(),
        timeout: Duration::from_secs(config.screen.timeout_secs),
        save_dir: None,
    });
    if !config.screen.enabled || screen.available().is_err() {
        eprintln!("no capturable screen; test skipped");
        return None;
    }
    Some(screen)
}

/// Width of a JPEG, read from its SOF header.
fn jpeg_width(bytes: &[u8]) -> Option<u32> {
    let mut i = 2;
    while i + 9 < bytes.len() {
        if bytes[i] != 0xFF {
            return None;
        }
        let marker = bytes[i + 1];
        let len = u16::from_be_bytes([bytes[i + 2], bytes[i + 3]]) as usize;
        if (0xC0..=0xCF).contains(&marker) && marker != 0xC4 && marker != 0xC8 && marker != 0xCC {
            return Some(u16::from_be_bytes([bytes[i + 7], bytes[i + 8]]) as u32);
        }
        i += 2 + len;
    }
    None
}

#[test]
fn camera_captures_a_jpeg() {
    let config = config();
    let Some(source) = camera_source(&config) else {
        return;
    };

    let started = Instant::now();
    let frame = source.capture().expect("the capture failed");
    eprintln!("{} KB frame in {:?}", frame.len() / 1024, started.elapsed());

    assert!(frame.len() > 1000, "the frame is suspiciously small");
    // JPEG header: if it is missing, ffmpeg returned something else and the
    // model would reject it without saying why.
    assert_eq!(&frame[..2], &[0xFF, 0xD8], "does not look like a JPEG");
}

#[test]
fn screen_is_captured_at_the_configured_width() {
    let config = config();
    let Some(source) = screen_source(&config) else {
        return;
    };

    let started = Instant::now();
    let frame = source.capture().expect("the capture failed");
    eprintln!(
        "{} KB capture in {:?}",
        frame.len() / 1024,
        started.elapsed()
    );
    assert_eq!(&frame[..2], &[0xFF, 0xD8], "does not look like a JPEG");

    // Width is what separates reading from making things up, so it is checked
    // that the script really scaled the image down and did not send it full size.
    let width = jpeg_width(&frame).expect("could not read the JPEG width");
    assert_eq!(
        width, config.screen.width,
        "the capture came out at {width} px and {} was requested",
        config.screen.width
    );
}

#[test]
#[ignore = "turns the camera on and runs a vision pass; slow"]
fn model_describes_what_the_camera_sees() {
    let _gpu = exclusive();
    let Some((config, llm, _)) = clients() else {
        return;
    };
    let Some(source) = camera_source(&config) else {
        return;
    };

    let tool = asist_tools::VisionTool::camera(Box::new(source), std::sync::Arc::new(llm));
    let started = Instant::now();
    let out = tool
        .call(&serde_json::json!({ "pregunta": "¿Qué se ve en la imagen?" }))
        .expect("the description failed");

    eprintln!("vision in {:?}: {out}", started.elapsed());
    assert!(
        out.split_whitespace().count() >= 3,
        "empty or minimal description: {out}"
    );
    assert!(
        started.elapsed() < Duration::from_secs(20),
        "took {:?}; check camera.width/height",
        started.elapsed()
    );
}

#[test]
#[ignore = "captures the screen and runs a vision pass; slow"]
fn model_describes_the_screen() {
    let _gpu = exclusive();
    let Some((config, llm, _)) = clients() else {
        return;
    };
    let Some(source) = screen_source(&config) else {
        return;
    };

    let tool = asist_tools::VisionTool::screen(Box::new(source), std::sync::Arc::new(llm));
    let started = Instant::now();
    let out = tool
        .call(&serde_json::json!({ "pregunta": "¿Qué hay en la pantalla?" }))
        .expect("the description failed");

    eprintln!("screen in {:?}: {out}", started.elapsed());
    assert!(
        out.split_whitespace().count() >= 3,
        "empty or minimal description: {out}"
    );
    // 7.6 s were measured at 1280 px. If this shoots up, either the GPU is busy
    // or someone raised screen.width without looking at the cost.
    assert!(
        started.elapsed() < Duration::from_secs(30),
        "took {:?}; check screen.width",
        started.elapsed()
    );
}

#[test]
#[ignore = "goes online with both search backends; slow"]
fn both_search_backends_answer_the_same() {
    // Compares what each one returns for the same questions. It does not claim
    // which is better (that depends on the question) but that both work, and it
    // leaves the numbers visible to choose with data.
    let config = config();
    let questions = [
        "capital de Australia",
        "qué tiempo hace hoy en Buenos Aires",
    ];

    let mut backends: Vec<(&str, asist_tools::SearchBackend)> = Vec::new();
    if let Ok(key) = std::env::var(&config.search.api_key_env) {
        if !key.trim().is_empty() {
            backends.push((
                "tavily",
                asist_tools::SearchBackend::Tavily { api_key: key },
            ));
        }
    }
    if let Some(program) = config.search.command.first() {
        if PathBuf::from(program).exists() {
            backends.push(("ddgs", asist_tools::SearchBackend::ddgs(program)));
        }
    }
    if backends.is_empty() {
        eprintln!("no search backend configured; test skipped");
        return;
    }

    for (name_in, backend) in backends {
        let tool = asist_tools::WebSearch::new(backend, 3, Duration::from_secs(30));
        for question in questions {
            let started = Instant::now();
            match tool.call(&serde_json::json!({ "consulta": question })) {
                Ok(out) => {
                    let first_answer = out.lines().next().unwrap_or("").to_string();
                    eprintln!(
                        "{name_in:7} {:>6.2}s  {:>5} bytes  «{question}»\n           {}",
                        started.elapsed().as_secs_f32(),
                        out.len(),
                        &first_answer[..first_answer.len().min(150)]
                    );
                    assert!(!out.trim().is_empty());
                }
                Err(err) => panic!("{name_in} failed on «{question}»: {err}"),
            }
        }
    }
}