//! Minimal OpenAI-compatible chat client (llama-server / LiteLLM proxy).

use anyhow::{Context, Result};
use base64::Engine as _;
use futures_util::StreamExt as _;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::Path;
use std::time::Duration;

use crate::config;
use crate::tools::{self, Policy, ToolCall};

#[derive(Serialize)]
struct ChatRequest<'a> {
    model: &'a str,
    messages: &'a [ChatMessage],
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<serde_json::Value>,
    max_tokens: u32,
    temperature: f32,
    stream: bool,
}

#[derive(Serialize, Deserialize, Clone)]
pub struct ChatMessage {
    pub role: String,
    /// Null for assistant turns that only carry tool calls.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub content: Option<Content>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_calls: Vec<ToolCall>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

impl ChatMessage {
    fn system(text: &str) -> Self {
        Self {
            role: "system".into(),
            content: Some(Content::Text(text.to_string())),
            tool_calls: Vec::new(),
            tool_call_id: None,
        }
    }

    fn user(content: Content) -> Self {
        Self {
            role: "user".into(),
            content: Some(content),
            tool_calls: Vec::new(),
            tool_call_id: None,
        }
    }

    fn assistant(content: Option<String>, tool_calls: Vec<ToolCall>) -> Self {
        Self {
            role: "assistant".into(),
            content: content.map(Content::Text),
            tool_calls,
            tool_call_id: None,
        }
    }

    fn tool_result(id: String, text: String) -> Self {
        Self {
            role: "tool".into(),
            content: Some(Content::Text(text)),
            tool_calls: Vec::new(),
            tool_call_id: Some(id),
        }
    }
}

/// Message content: plain text, or a list of text/image parts.
#[derive(Serialize, Deserialize, Clone)]
#[serde(untagged)]
pub enum Content {
    Text(String),
    Parts(Vec<Part>),
}

#[derive(Serialize, Deserialize, Clone)]
#[serde(tag = "type")]
pub enum Part {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "image_url")]
    ImageUrl { image_url: ImageUrl },
}

#[derive(Serialize, Deserialize, Clone)]
pub struct ImageUrl {
    pub url: String,
}

/// A local image to attach to the prompt.
pub struct Image {
    pub mime: &'static str,
    pub data: Vec<u8>,
}

impl Image {
    /// Guess the MIME type from the path's extension. Returns None when the
    /// extension isn't a supported image type, so callers can decide whether
    /// that's an error (one-shot argv) or just not an image (REPL input).
    pub fn mime_for(path: &Path) -> Option<&'static str> {
        match path
            .extension()
            .and_then(|e| e.to_str())
            .unwrap_or_default()
            .to_ascii_lowercase()
            .as_str()
        {
            "png" => Some("image/png"),
            "jpg" | "jpeg" => Some("image/jpeg"),
            "gif" => Some("image/gif"),
            "webp" => Some("image/webp"),
            _ => None,
        }
    }

    /// Read an image from disk, guessing the MIME type from the extension.
    pub fn load(path: &Path) -> Result<Self> {
        let mime = Self::mime_for(path).ok_or_else(|| {
            let ext = path
                .extension()
                .and_then(|e| e.to_str())
                .unwrap_or_default();
            anyhow::anyhow!(
                "unsupported image type '.{ext}' (supported: png, jpg, jpeg, gif, webp)"
            )
        })?;

        let data =
            std::fs::read(path).with_context(|| format!("reading image {}", path.display()))?;
        if data.is_empty() {
            anyhow::bail!("image {} is empty", path.display());
        }
        if data.len() > config::MAX_IMAGE_BYTES {
            anyhow::bail!(
                "image {} is {} MB, over the {} MB limit",
                path.display(),
                data.len() / 1_000_000,
                config::MAX_IMAGE_BYTES / 1_000_000
            );
        }
        Ok(Self { mime, data })
    }

    fn data_url(&self) -> String {
        let b64 = base64::engine::general_purpose::STANDARD.encode(&self.data);
        format!("data:{};base64,{b64}", self.mime)
    }
}

/// Streamed chunk of a completion.
#[derive(Deserialize)]
struct StreamChunk {
    #[serde(default)]
    choices: Vec<StreamChoice>,
}

#[derive(Deserialize)]
struct StreamChoice {
    #[serde(default)]
    delta: Option<StreamDelta>,
}

#[derive(Deserialize)]
struct StreamDelta {
    #[serde(default)]
    content: Option<Content>,
    #[serde(default)]
    tool_calls: Vec<ToolCallDelta>,
}

/// Tool calls arrive as fragments spread across many frames, keyed by `index`.
/// The name and id show up once; `arguments` is a JSON string built up a few
/// characters at a time.
#[derive(Deserialize)]
struct ToolCallDelta {
    #[serde(default)]
    index: usize,
    #[serde(default)]
    id: Option<String>,
    #[serde(default)]
    function: Option<FunctionDelta>,
}

#[derive(Deserialize)]
struct FunctionDelta {
    #[serde(default)]
    name: Option<String>,
    #[serde(default)]
    arguments: Option<String>,
}

/// Accumulates fragmented tool calls, keyed by their wire index.
#[derive(Default)]
struct ToolCallAccumulator {
    calls: BTreeMap<usize, ToolCall>,
}

impl ToolCallAccumulator {
    fn push(&mut self, deltas: Vec<ToolCallDelta>) {
        for delta in deltas {
            let call = self.calls.entry(delta.index).or_insert_with(|| ToolCall {
                id: String::new(),
                kind: "function".into(),
                function: crate::tools::FunctionCall {
                    name: String::new(),
                    arguments: String::new(),
                },
            });
            if let Some(id) = delta.id {
                call.id = id;
            }
            if let Some(function) = delta.function {
                if let Some(name) = function.name {
                    call.function.name = name;
                }
                if let Some(arguments) = function.arguments {
                    call.function.arguments.push_str(&arguments);
                }
            }
        }
    }

    fn finish(self) -> Vec<ToolCall> {
        let mut calls: Vec<ToolCall> = self.calls.into_values().collect();
        for (n, call) in calls.iter_mut().enumerate() {
            // Some servers omit ids; synthesize one so the tool reply still matches.
            if call.id.is_empty() {
                call.id = format!("call_{n}");
            }
        }
        calls
    }
}

/// An ongoing conversation: holds the message history and streams each reply.
///
/// One `send` may span several requests: if the model asks to call tools, the
/// calls run, their results go back as `tool` messages, and generation
/// continues until it produces a plain answer.
pub struct Chat {
    client: reqwest::Client,
    messages: Vec<ChatMessage>,
    policy: Policy,
}

impl Chat {
    pub fn new(policy: Policy) -> Result<Self> {
        let client = reqwest::Client::builder()
            .connect_timeout(Duration::from_secs(config::CONNECT_TIMEOUT_SECS))
            // Per-chunk idle timeout, not a total cap: long generations are
            // fine, but a stalled connection fails instead of hanging forever.
            .read_timeout(Duration::from_secs(config::READ_TIMEOUT_SECS))
            .build()
            .context("building HTTP client")?;

        let mut chat = Self {
            client,
            messages: Vec::new(),
            policy,
        };
        chat.reset();
        Ok(chat)
    }

    /// Drop the history and start over from the system prompt.
    pub fn reset(&mut self) {
        self.messages.clear();
        self.messages
            .push(ChatMessage::system(config::SYSTEM_PROMPT));
    }

    /// Number of turns so far, excluding the system prompt.
    pub fn turn_count(&self) -> usize {
        self.messages.len().saturating_sub(1)
    }

    /// Append a user message and stream the reply into `on_item`, which
    /// receives text deltas and tool events as one ordered stream.
    ///
    /// If a request fails, every message from this turn is rolled back so the
    /// history stays consistent with what the model actually saw.
    pub async fn send(
        &mut self,
        user_message: &str,
        images: &[Image],
        mut on_item: impl FnMut(Item),
    ) -> Result<()> {
        let mark = self.messages.len();
        self.messages
            .push(ChatMessage::user(user_content(user_message, images)));

        match self.run_turn(&mut on_item).await {
            Ok(()) => Ok(()),
            Err(err) => {
                self.messages.truncate(mark);
                Err(err)
            }
        }
    }

    async fn run_turn(&mut self, on_item: &mut impl FnMut(Item)) -> Result<()> {
        for round in 0..config::MAX_TOOL_ROUNDS {
            let (reply, tool_calls) = self.stream_reply(on_item).await?;

            if tool_calls.is_empty() {
                self.messages.push(ChatMessage::assistant(
                    (!reply.is_empty()).then_some(reply),
                    Vec::new(),
                ));
                return Ok(());
            }

            // Out of rounds: stop calling tools but still ask for an answer,
            // so everything gathered so far gets summarized rather than lost.
            if round + 1 == config::MAX_TOOL_ROUNDS {
                self.messages
                    .push(ChatMessage::assistant(Some(reply), tool_calls));
                self.messages.push(ChatMessage::user(Content::Text(
                    "Tool budget exhausted. Answer using only what you have gathered so far."
                        .into(),
                )));
                let (final_reply, _) = self.stream_reply(on_item).await?;
                self.messages.push(ChatMessage::assistant(
                    (!final_reply.is_empty()).then_some(final_reply),
                    Vec::new(),
                ));
                return Ok(());
            }

            // Record the assistant turn, tool calls included: the server needs
            // them to match the results that follow.
            self.messages
                .push(ChatMessage::assistant(Some(reply), tool_calls.clone()));

            for call in &tool_calls {
                on_item(Item::ToolStart {
                    name: call.function.name.clone(),
                    summary: call.summary(),
                });
                let result = tools::execute(call, self.policy, &mut blocking_confirm).await;
                on_item(Item::ToolDone {
                    outcome: outcome_summary(&result),
                });
                self.messages
                    .push(ChatMessage::tool_result(call.id.clone(), result));
            }
        }

        Ok(())
    }

    async fn stream_reply(
        &self,
        on_item: &mut impl FnMut(Item),
    ) -> Result<(String, Vec<ToolCall>)> {
        let body = ChatRequest {
            model: config::MODEL,
            messages: &self.messages,
            tools: Some(tools::specs()),
            max_tokens: config::MAX_TOKENS,
            temperature: config::TEMPERATURE,
            stream: true,
        };

        let url = format!("{}/v1/chat/completions", config::BASE_URL);
        let resp = self
            .client
            .post(&url)
            .bearer_auth(config::api_key())
            .json(&body)
            .send()
            .await
            .with_context(|| format!("request to {url}"))?;

        let status = resp.status();
        if !status.is_success() {
            let text = resp.text().await.context("reading error body")?;
            anyhow::bail!(
                "HTTP {status}: {}",
                text.chars().take(4000).collect::<String>()
            );
        }

        let mut stream = resp.bytes_stream();
        let mut reply = String::new();
        let mut buf = String::new();
        let mut accumulator = ToolCallAccumulator::default();

        'outer: while let Some(item) = stream.next().await {
            let bytes = item.context("reading stream chunk")?;
            buf.push_str(&String::from_utf8_lossy(&bytes));

            // SSE frames are line-delimited; the last one is incomplete until more arrives.
            while let Some(pos) = buf.find('\n') {
                let line: String = buf.drain(..=pos).collect();
                let line = line.trim();
                let Some(data) = line.strip_prefix("data:").map(str::trim) else {
                    continue; // comments, "event:" lines, blank lines
                };
                if data == "[DONE]" {
                    break 'outer;
                }
                let Ok(chunk) = serde_json::from_str::<StreamChunk>(data) else {
                    continue; // tolerate frames we don't understand
                };
                for choice in chunk.choices {
                    let Some(delta) = choice.delta else { continue };
                    if let Some(text) = delta_content(delta.content).filter(|t| !t.is_empty()) {
                        reply.push_str(&text);
                        on_item(Item::Text(text));
                    }
                    accumulator.push(delta.tool_calls);
                }
            }
        }

        Ok((reply, accumulator.finish()))
    }
}

/// Something surfaced to the UI mid-turn: a chunk of the model's answer, or a
/// note that a tool is running and how it turned out. One type so callers can
/// hand a single callback to `Chat::send`.
pub enum Item {
    Text(String),
    ToolStart { name: String, summary: String },
    ToolDone { outcome: String },
}

/// One line describing how a tool call went, for the status display.
fn outcome_summary(result: &str) -> String {
    let trimmed = result.trim_start();
    if let Some(rest) = trimmed.strip_prefix("Error:") {
        return format!(
            "error: {}",
            rest.trim().chars().take(70).collect::<String>()
        );
    }
    if trimmed.starts_with("Skipped:") {
        return "skipped".into();
    }
    let lines = result.lines().count();
    let first = result.trim().lines().next().unwrap_or("");
    let shown: String = first.chars().take(48).collect();
    format!("{lines} line(s) — {shown}")
}

/// Confirmation has to block on terminal input. That's safe here: it only
/// happens while a tool is running, with no other request in flight, so
/// parking the task costs nothing.
///
/// Defaults to yes: a bare Enter proceeds, `n` declines.
fn blocking_confirm(_call: &ToolCall) -> bool {
    use std::io::Write as _;
    eprint!("  run it? [Y/n] ");
    let _ = std::io::stderr().flush();

    let mut answer = String::new();
    if std::io::stdin().read_line(&mut answer).unwrap_or(0) == 0 {
        return true; // EOF, nothing left to ask
    }
    !matches!(answer.trim(), "n" | "N" | "no" | "No")
}

/// Build the content for a user turn: plain text, or text plus image parts.
/// Empty text with images still yields image-only parts.
fn user_content(user_message: &str, images: &[Image]) -> Content {
    if images.is_empty() {
        return Content::Text(user_message.to_string());
    }

    let mut parts = Vec::with_capacity(images.len() + 1);
    if !user_message.is_empty() {
        parts.push(Part::Text {
            text: user_message.to_string(),
        });
    }
    for image in images {
        parts.push(Part::ImageUrl {
            image_url: ImageUrl {
                url: image.data_url(),
            },
        });
    }
    Content::Parts(parts)
}

/// Delta content is a string on every server we target, but some emit content
/// parts; flatten either shape.
fn delta_content(content: Option<Content>) -> Option<String> {
    match content? {
        Content::Text(text) => Some(text),
        Content::Parts(parts) => {
            let mut out = String::new();
            for part in parts {
                if let Part::Text { text } = part {
                    out.push_str(&text);
                }
            }
            Some(out)
        }
    }
}