//! Transport to llama-server.

use std::time::{Duration, Instant};

use serde_json::{json, Value};

use asist_core::config::LlmConfig;
use asist_core::error::{Error, Result};
use asist_core::http::{Cancel, HttpClient};
use asist_core::tools::{ToolCall, ToolRegistry};

use crate::chat::Conversation;

/// A piece of the answer as it comes out of the model.
#[derive(Debug, Clone)]
pub enum Delta {
    /// Text for the user.
    Text(String),
    /// The model decided to call tools; no more text will come.
    ToolCalls(Vec<ToolCall>),
}

/// How a streaming answer ended.
#[derive(Debug, Clone, Default)]
pub struct StreamOutcome {
    pub text: String,
    pub tool_calls: Vec<ToolCall>,
    /// Time until the first fragment with content.
    pub ttft: Option<Duration>,
    pub stopped_early: bool,
}

pub struct LlmClient {
    http: HttpClient,
    config: LlmConfig,
}

impl LlmClient {
    pub fn new(authority: String, config: &LlmConfig) -> Self {
        Self {
            http: HttpClient::new(authority)
                .with_read_timeout(Duration::from_secs(config.request_timeout_secs)),
            config: config.clone(),
        }
    }

    pub fn healthy(&self) -> bool {
        self.http.healthy("/health")
    }

    /// Waits for the server to finish loading the model.
    pub fn wait_ready(&self, timeout: Duration) -> Result<()> {
        let deadline = Instant::now() + timeout;
        while Instant::now() < deadline {
            if self.healthy() {
                return Ok(());
            }
            std::thread::sleep(Duration::from_millis(500));
        }
        Err(Error::Llm(format!(
            "{} did not answer /health within {} s",
            self.http.authority,
            timeout.as_secs()
        )))
    }

    /// Checks that the chat template does not open a reasoning block
    /// without closing it.
    ///
    /// It is worth a warning: with the original Qwen3.5 template the model
    /// spends 7 to 9 seconds «thinking» before the first audible word, and from
    /// the outside the assistant looks hung.
    pub fn warn_if_thinking_template(&self) {
        let Ok(body) = self.http.get("/props") else {
            return;
        };
        let Ok(props) = serde_json::from_slice::<Value>(&body) else {
            return;
        };
        let Some(template) = props.get("chat_template").and_then(Value::as_str) else {
            return;
        };

        let opens = template.matches("<think>").count();
        let closes = template.matches("</think>").count();
        if opens > closes {
            tracing::warn!(
                target: "llm",
                "the chat template opens <think> without closing it: the model will reason \
                 for several seconds before answering. Start llama-server with \
                 --chat-template-file config/qwen35-no-think.jinja"
            );
        }
    }

    /// Asks the model about an image, in a request separate from the
    /// conversation.
    ///
    /// The server has the multimodal projector loaded (`--mmproj`), so it
    /// accepts `image_url` parts with the image in base64. It is kept out of the
    /// history on purpose: an image takes hundreds of context tokens and dragging
    /// it turn after turn would be very expensive for how little it adds once
    /// described. What goes back into the conversation is the text.
    ///
    /// Measured on this machine, the cost depends heavily on size: 1.3 s at
    /// 320x240, 2.9 s at 640x480 and 7.8 s at 1280x720.
    pub fn look(&self, image_jpeg: &[u8], question: &str, cancel: &Cancel) -> Result<String> {
        use base64::Engine;
        let encoded = base64::engine::general_purpose::STANDARD.encode(image_jpeg);

        let mut body = json!({
            "messages": [{
                "role": "user",
                "content": [
                    { "type": "text", "text": question },
                    {
                        "type": "image_url",
                        "image_url": { "url": format!("data:image/jpeg;base64,{encoded}") }
                    }
                ]
            }],
            "stream": true,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
        });
        if !self.config.model.is_empty() {
            body["model"] = json!(self.config.model);
        }

        // Streaming even though the partial text is unused: a vision request
        // takes seconds and without a stream the socket may stay silent past the
        // read deadline.
        let started = Instant::now();
        let mut text = String::new();
        let mut finished = false;
        let result = self
            .http
            .post_json_lines("/v1/chat/completions", &body, cancel, |line| {
                let Some(payload) = line.strip_prefix("data: ") else {
                    return true;
                };
                if payload.trim() == "[DONE]" {
                    finished = true;
                    return false;
                }
                let Ok(event) = serde_json::from_str::<Value>(payload) else {
                    return true;
                };
                let Some(choice) = event.get("choices").and_then(|c| c.get(0)) else {
                    return true;
                };
                if let Some(chunk) = choice
                    .get("delta")
                    .and_then(|d| d.get("content"))
                    .and_then(Value::as_str)
                {
                    text.push_str(chunk);
                }
                if choice.get("finish_reason").is_some_and(|r| !r.is_null()) {
                    finished = true;
                    return false;
                }
                true
            });
        match result {
            Ok(()) => {}
            Err(Error::Cancelled) if finished || cancel.is_cancelled() => {}
            Err(err) => return Err(err),
        }

        tracing::debug!(
            target: "llm",
            bytes = image_jpeg.len(),
            ms = started.elapsed().as_millis(),
            "vision"
        );
        let text = text.trim().to_string();
        if text.is_empty() {
            return Err(Error::Llm("the model did not describe the image".into()));
        }
        Ok(text)
    }

    fn request_body(&self, chat: &Conversation, tools: Option<&ToolRegistry>) -> Value {
        let mut body = json!({
            "messages": chat.to_json(),
            "stream": true,
            "temperature": self.config.temperature,
            "top_p": self.config.top_p,
            "top_k": self.config.top_k,
            "min_p": self.config.min_p,
            "repeat_penalty": self.config.repeat_penalty,
            "max_tokens": self.config.max_tokens,
        });
        if !self.config.model.is_empty() {
            body["model"] = json!(self.config.model);
        }
        if let Some(tools) = tools.filter(|t| !t.is_empty()) {
            body["tools"] = tools.schema();
            body["tool_choice"] = json!("auto");
        }
        body
    }

    /// Sends the conversation and delivers the answer in pieces.
    ///
    /// `on_delta` returns `false` to stop (what the orchestrator does when the
    /// user interrupts); `cancel` does the same from another thread.
    pub fn stream(
        &self,
        chat: &Conversation,
        tools: Option<&ToolRegistry>,
        cancel: &Cancel,
        mut on_delta: impl FnMut(&Delta) -> bool,
    ) -> Result<StreamOutcome> {
        let body = self.request_body(chat, tools);
        let started = Instant::now();
        let mut outcome = StreamOutcome::default();
        let mut assembler = ToolCallAssembler::default();
        let mut finished = false;

        let result = self
            .http
            .post_json_lines("/v1/chat/completions", &body, cancel, |line| {
                let Some(payload) = line.strip_prefix("data: ") else {
                    return true;
                };
                if payload.trim() == "[DONE]" {
                    finished = true;
                    return false;
                }
                let Ok(event) = serde_json::from_str::<Value>(payload) else {
                    tracing::debug!(target: "llm", %payload, "fragmento SSE ilegible");
                    return true;
                };
                let Some(choice) = event.get("choices").and_then(|c| c.get(0)) else {
                    return true;
                };
                let delta = choice.get("delta").unwrap_or(&Value::Null);

                if let Some(calls) = delta.get("tool_calls").and_then(Value::as_array) {
                    assembler.absorb(calls);
                }
                if let Some(text) = delta.get("content").and_then(Value::as_str) {
                    if !text.is_empty() {
                        if outcome.ttft.is_none() {
                            outcome.ttft = Some(started.elapsed());
                        }
                        outcome.text.push_str(text);
                        if !on_delta(&Delta::Text(text.to_string())) {
                            outcome.stopped_early = true;
                            return false;
                        }
                    }
                }
                // A `finish_reason` closes the answer even if the server never
                // sends [DONE] (it happens when the connection is cut).
                if choice.get("finish_reason").is_some_and(|r| !r.is_null()) {
                    finished = true;
                    return false;
                }
                true
            });

        match result {
            Ok(()) => {}
            // Stopping on purpose is not a failure: `on_delta` returned false or
            // cancellation was triggered.
            Err(Error::Cancelled) if finished || outcome.stopped_early || cancel.is_cancelled() => {
            }
            Err(err) => return Err(err),
        }

        outcome.tool_calls = assembler.finish();
        if !outcome.tool_calls.is_empty()
            && !on_delta(&Delta::ToolCalls(outcome.tool_calls.clone()))
        {
            outcome.stopped_early = true;
        }
        Ok(outcome)
    }
}

/// Reassembles tool calls, which arrive split across fragments.
///
/// Streaming sends the name in one fragment and the arguments in several
/// more, identified only by their `index`; until the stream ends there is
/// no complete call to run.
#[derive(Default)]
struct ToolCallAssembler {
    partial: Vec<PartialCall>,
}

#[derive(Default, Clone)]
struct PartialCall {
    id: String,
    name: String,
    arguments: String,
}

impl ToolCallAssembler {
    fn absorb(&mut self, calls: &[Value]) {
        for call in calls {
            let index = call.get("index").and_then(Value::as_u64).unwrap_or(0) as usize;
            if self.partial.len() <= index {
                self.partial.resize(index + 1, PartialCall::default());
            }
            let slot = &mut self.partial[index];
            if let Some(id) = call.get("id").and_then(Value::as_str) {
                if !id.is_empty() {
                    slot.id = id.to_string();
                }
            }
            let Some(function) = call.get("function") else {
                continue;
            };
            if let Some(name) = function.get("name").and_then(Value::as_str) {
                if !name.is_empty() {
                    slot.name = name.to_string();
                }
            }
            if let Some(args) = function.get("arguments").and_then(Value::as_str) {
                slot.arguments.push_str(args);
            }
        }
    }

    fn finish(self) -> Vec<ToolCall> {
        self.partial
            .into_iter()
            .enumerate()
            .filter(|(_, call)| !call.name.is_empty())
            .map(|(i, call)| ToolCall {
                id: if call.id.is_empty() {
                    format!("call_{i}")
                } else {
                    call.id
                },
                name: call.name,
                arguments: if call.arguments.is_empty() {
                    "{}".into()
                } else {
                    call.arguments
                },
            })
            .collect()
    }
}

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

    fn absorb(assembler: &mut ToolCallAssembler, raw: &str) {
        let value: Value = serde_json::from_str(raw).unwrap();
        assembler.absorb(value.as_array().unwrap());
    }

    #[test]
    fn chunked_arguments_are_reassembled() {
        let mut assembler = ToolCallAssembler::default();
        absorb(
            &mut assembler,
            r#"[{"index":0,"id":"c1","function":{"name":"eco","arguments":"{\"a\":"}}]"#,
        );
        absorb(
            &mut assembler,
            r#"[{"index":0,"function":{"arguments":"1}"}}]"#,
        );
        let calls = assembler.finish();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].id, "c1");
        assert_eq!(calls[0].name, "eco");
        assert_eq!(calls[0].arguments, r#"{"a":1}"#);
    }

    #[test]
    fn several_parallel_calls_are_reassembled() {
        let mut assembler = ToolCallAssembler::default();
        absorb(
            &mut assembler,
            r#"[{"index":0,"id":"a","function":{"name":"uno","arguments":"{}"}}]"#,
        );
        absorb(
            &mut assembler,
            r#"[{"index":1,"id":"b","function":{"name":"dos","arguments":"{}"}}]"#,
        );
        let calls = assembler.finish();
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].name, "uno");
        assert_eq!(calls[1].name, "dos");
    }

    #[test]
    fn no_arguments_means_an_empty_object() {
        let mut assembler = ToolCallAssembler::default();
        absorb(
            &mut assembler,
            r#"[{"index":0,"id":"c","function":{"name":"hora_actual"}}]"#,
        );
        assert_eq!(assembler.finish()[0].arguments, "{}");
    }

    #[test]
    fn slot_without_a_name_produces_no_call() {
        // The server may number index 1 without ever having sent 0.
        let mut assembler = ToolCallAssembler::default();
        absorb(
            &mut assembler,
            r#"[{"index":1,"id":"b","function":{"name":"dos","arguments":"{}"}}]"#,
        );
        let calls = assembler.finish();
        assert_eq!(calls.len(), 1, "the gap must not become an empty call");
        assert_eq!(calls[0].name, "dos");
    }

    #[test]
    fn without_an_id_a_stable_one_is_generated() {
        let mut assembler = ToolCallAssembler::default();
        absorb(
            &mut assembler,
            r#"[{"index":0,"function":{"name":"x","arguments":"{}"}}]"#,
        );
        assert_eq!(assembler.finish()[0].id, "call_0");
    }
}