//! The llama inference server's chat-completions types and the request builder.

pub mod assemble;
pub mod info;
pub mod request;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChatMessage {
    System {
        content: String,
    },
    User {
        content: String,
    },
    Assistant {
        content: Option<String>,
        reasoning_content: Option<String>,
        tool_calls: Vec<proto::ToolCall>,
    },
    Tool {
        tool_call_id: String,
        content: String,
    },
}

#[derive(Debug, Clone, PartialEq)]
pub struct ToolSchema {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

#[derive(Debug, Clone, PartialEq)]
pub struct ChatRequest {
    pub slot: u32,
    pub messages: Vec<ChatMessage>,
    pub tools: Vec<ToolSchema>,
    pub thinking: bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChatEvent {
    Queued {
        ahead: usize,
    },
    Waiting {
        slot_busy: bool,
    },
    Progress {
        total: u64,
        cache: u64,
        processed: u64,
    },
    Reasoning(String),
    Content(String),
    ToolCallDelta {
        index: u32,
        id: Option<String>,
        name: Option<String>,
        arguments: String,
    },
    ThinkingCapped {
        tokens: u64,
    },
    Retrying {
        attempt: u32,
        after_ms: u64,
        error: String,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FinishReason {
    Stop,
    ToolCalls,
    Length,
}

/// The three counts loopd uses. The server sends more fields beside them; they are ignored.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize)]
pub struct Timings {
    pub cache_n: u64,
    pub prompt_n: u64,
    pub predicted_n: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Completion {
    pub id: String,
    pub content: Option<String>,
    pub reasoning_content: Option<String>,
    pub tool_calls: Vec<proto::ToolCall>,
    pub finish_reason: FinishReason,
    pub timings: Timings,
    pub reasoning_tokens: u64,
    pub thinking_capped: bool,
}

#[derive(Debug)]
pub enum InferError {
    Busy,
    Connect(std::io::Error),
    WaitTimeout,
    LoadTimeout,
    Stalled,
    StreamClosedEarly,
    ThinkingOverrun,
    Http { status: u16, body: String },
    Protocol(String),
}

impl std::fmt::Display for InferError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            InferError::Busy => write!(f, "the slot is busy"),
            InferError::Connect(err) => {
                write!(f, "could not connect to the inference server: {err}")
            }
            InferError::WaitTimeout => write!(f, "timed out waiting for the slot to free up"),
            InferError::LoadTimeout => write!(f, "timed out waiting for the model to load"),
            InferError::Stalled => write!(f, "the stream stalled"),
            InferError::StreamClosedEarly => {
                write!(f, "the stream closed before the response was complete")
            }
            InferError::ThinkingOverrun => write!(f, "thinking ran over its token budget"),
            InferError::Http { status, body } => {
                write!(f, "the server responded with {status}: {body}")
            }
            InferError::Protocol(what) => write!(f, "the server spoke the protocol wrong: {what}"),
        }
    }
}

impl std::error::Error for InferError {}

pub struct Client {
    pub(crate) cfg: crate::config::Config,
}

impl Client {
    pub fn new(cfg: crate::config::Config) -> Self {
        Self { cfg }
    }
    pub fn config(&self) -> &crate::config::Config {
        &self.cfg
    }
}