//! Provider-neutral exact-endpoint LLM transport and bounded wire envelope.

#![allow(clippy::missing_errors_doc)] // Public methods share the exhaustive typed LlmError boundary.

use crate::config::LlmConnection;
use crate::types::{
    BoundedText, BoundedVec, MAX_BODY_BYTES, MAX_CONVERSATION_MESSAGES, MAX_IDENTIFIER_BYTES,
    MAX_MESSAGE_BYTES, MAX_TOOL_CALLS, MessageRole, ProposedToolCall,
};
use libremetaverse_types::compat::CancellationToken;
use serde::Deserialize;
use serde_json::{Map, Value, json};
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;

const MAX_TOOLS: usize = 64;
const MAX_SCHEMA_PROPERTIES: usize = 128;
const MAX_SCHEMA_DEPTH: usize = 16;
const MAX_SCHEMA_NODES: usize = 1_024;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ContentPart {
    Text(BoundedText<MAX_MESSAGE_BYTES>),
    Image {
        url: BoundedText<MAX_BODY_BYTES>,
        detail: ImageDetail,
    },
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ImageDetail {
    Auto,
    Low,
    High,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CompletionMessage {
    pub role: MessageRole,
    pub content: BoundedVec<ContentPart, 16>,
    pub tool_call_id: Option<BoundedText<MAX_IDENTIFIER_BYTES>>,
    pub proposed_calls: BoundedVec<ProposedToolCall, MAX_TOOL_CALLS>,
}

impl CompletionMessage {
    pub fn text(role: MessageRole, text: impl Into<String>) -> Result<Self, LlmError> {
        let mut content = BoundedVec::new();
        content
            .try_push(
                "completion_message.content",
                ContentPart::Text(BoundedText::new("completion_message.text", text)?),
            )
            .map_err(LlmError::Boundary)?;
        Ok(Self {
            role,
            content,
            tool_call_id: None,
            proposed_calls: BoundedVec::new(),
        })
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ToolSchema {
    Object {
        properties: BTreeMap<String, ToolSchema>,
        required: BTreeSet<String>,
        additional_properties: bool,
    },
    String,
    Integer,
    Number,
    Boolean,
    Array {
        items: Box<ToolSchema>,
        max_items: usize,
    },
}

impl ToolSchema {
    pub fn validate_schema(&self) -> Result<(), LlmError> {
        let mut remaining_nodes = MAX_SCHEMA_NODES;
        self.validate_schema_inner(0, &mut remaining_nodes)
    }

    fn validate_schema_inner(
        &self,
        depth: usize,
        remaining_nodes: &mut usize,
    ) -> Result<(), LlmError> {
        if depth > MAX_SCHEMA_DEPTH || *remaining_nodes == 0 {
            return Err(LlmError::InvalidToolSchema);
        }
        *remaining_nodes -= 1;
        match self {
            Self::Object {
                properties,
                required,
                ..
            } => {
                if properties.len() > MAX_SCHEMA_PROPERTIES
                    || required.len() > MAX_SCHEMA_PROPERTIES
                    || required.iter().any(|name| !properties.contains_key(name))
                    || properties
                        .keys()
                        .any(|name| name.is_empty() || name.len() > MAX_IDENTIFIER_BYTES)
                {
                    return Err(LlmError::InvalidToolSchema);
                }
                for schema in properties.values() {
                    schema.validate_schema_inner(depth + 1, remaining_nodes)?;
                }
                Ok(())
            }
            Self::Array { items, max_items } => {
                if *max_items == 0 || *max_items > 256 {
                    return Err(LlmError::InvalidToolSchema);
                }
                items.validate_schema_inner(depth + 1, remaining_nodes)
            }
            _ => Ok(()),
        }
    }

    pub fn validate_value(&self, value: &Value) -> Result<(), LlmError> {
        match self {
            Self::Object {
                properties,
                required,
                additional_properties,
            } => {
                let object = value.as_object().ok_or(LlmError::InvalidToolArguments)?;
                if required.iter().any(|name| !object.contains_key(name))
                    || (!additional_properties
                        && object.keys().any(|name| !properties.contains_key(name)))
                {
                    return Err(LlmError::InvalidToolArguments);
                }
                for (name, value) in object {
                    if let Some(schema) = properties.get(name) {
                        schema.validate_value(value)?;
                    }
                }
                Ok(())
            }
            Self::String if value.is_string() => Ok(()),
            Self::Integer if value.as_i64().is_some() || value.as_u64().is_some() => Ok(()),
            Self::Number if value.is_number() => Ok(()),
            Self::Boolean if value.is_boolean() => Ok(()),
            Self::Array { items, max_items } => {
                let values = value.as_array().ok_or(LlmError::InvalidToolArguments)?;
                if values.len() > *max_items {
                    return Err(LlmError::InvalidToolArguments);
                }
                for value in values {
                    items.validate_value(value)?;
                }
                Ok(())
            }
            _ => Err(LlmError::InvalidToolArguments),
        }
    }

    fn wire_value(&self) -> Value {
        match self {
            Self::Object {
                properties,
                required,
                additional_properties,
            } => {
                let properties = properties
                    .iter()
                    .map(|(name, schema)| (name.clone(), schema.wire_value()))
                    .collect::<Map<_, _>>();
                json!({
                    "type": "object",
                    "properties": properties,
                    "required": required,
                    "additionalProperties": additional_properties
                })
            }
            Self::String => json!({"type":"string"}),
            Self::Integer => json!({"type":"integer"}),
            Self::Number => json!({"type":"number"}),
            Self::Boolean => json!({"type":"boolean"}),
            Self::Array { items, max_items } => {
                json!({"type":"array", "items":items.wire_value(), "maxItems":max_items})
            }
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolDefinition {
    pub name: BoundedText<MAX_IDENTIFIER_BYTES>,
    pub description: BoundedText<MAX_MESSAGE_BYTES>,
    pub schema: ToolSchema,
    pub mutating: bool,
}

impl ToolDefinition {
    pub fn validate(&self) -> Result<(), LlmError> {
        if !self
            .name
            .chars()
            .all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '-'))
        {
            return Err(LlmError::InvalidToolSchema);
        }
        self.schema.validate_schema()
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Usage {
    pub prompt_tokens: Option<u64>,
    pub completion_tokens: Option<u64>,
    pub total_tokens: Option<u64>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Completion {
    pub request_id: u64,
    pub correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
    pub message: CompletionMessage,
    pub usage: Option<Usage>,
    pub latency: Duration,
    pub attempts: usize,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct LlmTransportLimits {
    pub connect_timeout: Duration,
    pub request_timeout: Duration,
    pub read_idle_timeout: Duration,
    pub pool_idle_timeout: Duration,
    pub total_timeout: Duration,
    pub max_prompt_bytes: usize,
    pub max_response_bytes: usize,
    pub max_concurrent_requests: usize,
    pub max_retries: usize,
    pub max_retry_delay: Duration,
}

impl Default for LlmTransportLimits {
    fn default() -> Self {
        Self {
            connect_timeout: Duration::from_secs(10),
            request_timeout: Duration::from_mins(1),
            read_idle_timeout: Duration::from_secs(15),
            pool_idle_timeout: Duration::from_secs(30),
            total_timeout: Duration::from_secs(90),
            max_prompt_bytes: 1024 * 1024,
            max_response_bytes: 1024 * 1024,
            max_concurrent_requests: 8,
            max_retries: 2,
            max_retry_delay: Duration::from_secs(5),
        }
    }
}

impl LlmTransportLimits {
    pub fn validate(&self) -> Result<(), LlmError> {
        if self.connect_timeout.is_zero()
            || self.request_timeout.is_zero()
            || self.read_idle_timeout.is_zero()
            || self.pool_idle_timeout.is_zero()
            || self.total_timeout.is_zero()
            || self.max_prompt_bytes == 0
            || self.max_prompt_bytes > MAX_BODY_BYTES
            || self.max_response_bytes == 0
            || self.max_response_bytes > MAX_BODY_BYTES
            || self.max_concurrent_requests == 0
            || self.max_concurrent_requests > 256
            || self.max_retries > 8
            || self.max_retry_delay > Duration::from_mins(1)
        {
            return Err(LlmError::UnsafeLimits);
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LlmError {
    Boundary(crate::types::BoundaryError),
    UnsafeLimits,
    InvalidToolSchema,
    InvalidToolArguments,
    PromptTooLarge,
    ResponseTooLarge,
    Cancelled,
    Timeout,
    Transport,
    RedirectRefused,
    HttpStatus(u16),
    MalformedJson,
    UnsupportedResponse,
    DuplicateToolCallId,
}

impl fmt::Display for LlmError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Boundary(error) => write!(formatter, "LLM boundary rejected input: {error}"),
            Self::UnsafeLimits => formatter.write_str("unsafe LLM transport limits"),
            Self::InvalidToolSchema => formatter.write_str("invalid registered tool schema"),
            Self::InvalidToolArguments => formatter.write_str("tool arguments do not match schema"),
            Self::PromptTooLarge => formatter.write_str("LLM prompt exceeds its byte bound"),
            Self::ResponseTooLarge => formatter.write_str("LLM response exceeds its byte bound"),
            Self::Cancelled => formatter.write_str("LLM request cancelled"),
            Self::Timeout => formatter.write_str("LLM request timed out"),
            Self::Transport => formatter.write_str("LLM transport failed"),
            Self::RedirectRefused => formatter.write_str("LLM endpoint redirect refused"),
            Self::HttpStatus(status) => write!(formatter, "LLM endpoint returned HTTP {status}"),
            Self::MalformedJson => formatter.write_str("LLM endpoint returned malformed JSON"),
            Self::UnsupportedResponse => formatter.write_str("unsupported LLM response shape"),
            Self::DuplicateToolCallId => {
                formatter.write_str("LLM response repeated a tool-call ID")
            }
        }
    }
}

impl Error for LlmError {}

impl From<crate::types::BoundaryError> for LlmError {
    fn from(value: crate::types::BoundaryError) -> Self {
        Self::Boundary(value)
    }
}

pub struct LlmClient {
    connection: LlmConnection,
    client: reqwest::Client,
    limits: LlmTransportLimits,
    slots: Arc<Semaphore>,
    next_request_id: AtomicU64,
}

impl fmt::Debug for LlmClient {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("LlmClient")
            .field("endpoint", &self.connection.endpoint_url)
            .field("limits", &self.limits)
            .field("available_slots", &self.slots.available_permits())
            .finish_non_exhaustive()
    }
}

impl LlmClient {
    pub fn new(connection: LlmConnection, limits: LlmTransportLimits) -> Result<Self, LlmError> {
        limits.validate()?;
        let client = reqwest::Client::builder()
            .redirect(reqwest::redirect::Policy::none())
            .connect_timeout(limits.connect_timeout)
            .timeout(limits.request_timeout)
            .read_timeout(limits.read_idle_timeout)
            .pool_idle_timeout(limits.pool_idle_timeout)
            .pool_max_idle_per_host(limits.max_concurrent_requests)
            .build()
            .map_err(|_| LlmError::Transport)?;
        Ok(Self {
            connection,
            client,
            slots: Arc::new(Semaphore::new(limits.max_concurrent_requests)),
            limits,
            next_request_id: AtomicU64::new(1),
        })
    }

    pub async fn complete(
        &self,
        messages: &[CompletionMessage],
        tools: &[ToolDefinition],
        cancellation: &CancellationToken,
    ) -> Result<Completion, LlmError> {
        if messages.is_empty()
            || messages.len() > MAX_CONVERSATION_MESSAGES
            || tools.len() > MAX_TOOLS
        {
            return Err(LlmError::PromptTooLarge);
        }
        for tool in tools {
            tool.validate()?;
        }
        let body = request_body(messages, tools)?;
        if body.len() > self.limits.max_prompt_bytes {
            return Err(LlmError::PromptTooLarge);
        }
        let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
        let correlation_id = BoundedText::new(
            "llm.correlation_id",
            format!("agent-request-{request_id:016x}"),
        )?;
        let started = Instant::now();
        let operation = async {
            let permit = tokio::select! {
                () = cancellation.cancelled() => return Err(LlmError::Cancelled),
                permit = Arc::clone(&self.slots).acquire_owned() => permit.map_err(|_| LlmError::Cancelled)?,
            };
            let result = self
                .complete_with_retries(request_id, correlation_id, body, cancellation, started)
                .await;
            drop(permit);
            result
        };
        tokio::time::timeout(self.limits.total_timeout, operation)
            .await
            .map_err(|_| LlmError::Timeout)?
    }

    async fn complete_with_retries(
        &self,
        request_id: u64,
        correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
        body: Vec<u8>,
        cancellation: &CancellationToken,
        started: Instant,
    ) -> Result<Completion, LlmError> {
        for attempt in 0..=self.limits.max_retries {
            match self
                .complete_once(
                    request_id,
                    correlation_id.clone(),
                    &body,
                    cancellation,
                    started,
                    attempt + 1,
                )
                .await
            {
                Ok(completion) => return Ok(completion),
                Err(AttemptError { error, retry_after })
                    if attempt < self.limits.max_retries && transient(&error) =>
                {
                    let delay = retry_after
                        .unwrap_or_else(|| deterministic_backoff(request_id, attempt))
                        .min(self.limits.max_retry_delay);
                    tokio::select! {
                        () = cancellation.cancelled() => return Err(LlmError::Cancelled),
                        () = tokio::time::sleep(delay) => {}
                    }
                }
                Err(AttemptError { error, .. }) => return Err(error),
            }
        }
        Err(LlmError::Transport)
    }

    async fn complete_once(
        &self,
        request_id: u64,
        correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
        body: &[u8],
        cancellation: &CancellationToken,
        started: Instant,
        attempts: usize,
    ) -> Result<Completion, AttemptError> {
        let response = tokio::select! {
            () = cancellation.cancelled() => return Err(AttemptError::new(LlmError::Cancelled)),
            response = self.client
                .post(self.connection.endpoint_url.expose_url())
                .bearer_auth(self.connection.api_key.expose_secret())
                .header("content-type", "application/json")
                .header("accept", "application/json")
                .header("x-correlation-id", correlation_id.as_str())
                .body(body.to_vec())
                .send() => response.map_err(|error| AttemptError::new(classify_reqwest(&error)))?,
        };
        let status = response.status();
        if status.is_redirection() {
            return Err(AttemptError::new(LlmError::RedirectRefused));
        }
        if !status.is_success() {
            let retry_after = response
                .headers()
                .get("retry-after")
                .and_then(|value| value.to_str().ok())
                .and_then(|value| value.parse::<u64>().ok())
                .map(Duration::from_secs);
            return Err(AttemptError {
                error: LlmError::HttpStatus(status.as_u16()),
                retry_after,
            });
        }
        let bytes = read_bounded(response, self.limits.max_response_bytes, cancellation).await?;
        parse_completion(
            request_id,
            correlation_id,
            &bytes,
            started.elapsed(),
            attempts,
        )
        .map_err(AttemptError::new)
    }
}

struct AttemptError {
    error: LlmError,
    retry_after: Option<Duration>,
}

impl AttemptError {
    const fn new(error: LlmError) -> Self {
        Self {
            error,
            retry_after: None,
        }
    }
}

async fn read_bounded(
    mut response: reqwest::Response,
    maximum: usize,
    cancellation: &CancellationToken,
) -> Result<Vec<u8>, AttemptError> {
    if response
        .content_length()
        .is_some_and(|length| usize::try_from(length).map_or(true, |length| length > maximum))
    {
        return Err(AttemptError::new(LlmError::ResponseTooLarge));
    }
    let mut body = Vec::with_capacity(
        response
            .content_length()
            .and_then(|value| usize::try_from(value).ok())
            .unwrap_or(0)
            .min(maximum),
    );
    loop {
        let chunk = tokio::select! {
            () = cancellation.cancelled() => return Err(AttemptError::new(LlmError::Cancelled)),
            chunk = response.chunk() => chunk.map_err(|error| AttemptError::new(classify_reqwest(&error)))?,
        };
        let Some(chunk) = chunk else { break };
        if body.len().saturating_add(chunk.len()) > maximum {
            return Err(AttemptError::new(LlmError::ResponseTooLarge));
        }
        body.extend_from_slice(&chunk);
    }
    Ok(body)
}

fn request_body(
    messages: &[CompletionMessage],
    tools: &[ToolDefinition],
) -> Result<Vec<u8>, LlmError> {
    let messages = messages.iter().map(message_wire_value).collect::<Vec<_>>();
    let tools = tools
        .iter()
        .map(|tool| {
            json!({
                "type":"function",
                "function": {
                    "name": tool.name.as_str(),
                    "description": tool.description.as_str(),
                    "parameters": tool.schema.wire_value()
                }
            })
        })
        .collect::<Vec<_>>();
    serde_json::to_vec(&json!({"messages":messages, "tools":tools}))
        .map_err(|_| LlmError::MalformedJson)
}

fn message_wire_value(message: &CompletionMessage) -> Value {
    let role = match message.role {
        MessageRole::System => "system",
        MessageRole::Avatar => "user",
        MessageRole::Agent => "assistant",
        MessageRole::Tool => "tool",
    };
    let content = message.content.as_slice().iter().map(|part| match part {
        ContentPart::Text(text) => json!({"type":"text", "text":text.as_str()}),
        ContentPart::Image { url, detail } => json!({
            "type":"image_url",
            "image_url":{"url":url.as_str(), "detail":match detail { ImageDetail::Auto=>"auto", ImageDetail::Low=>"low", ImageDetail::High=>"high"}}
        }),
    }).collect::<Vec<_>>();
    let mut value = json!({"role":role,"content":content});
    if let Some(call_id) = &message.tool_call_id {
        value["tool_call_id"] = Value::String(call_id.as_str().to_owned());
    }
    if !message.proposed_calls.is_empty() {
        value["tool_calls"] = Value::Array(message.proposed_calls.as_slice().iter().map(|call| json!({
            "id":call.call_id.as_str(), "type":"function", "function":{"name":call.name.as_str(),"arguments":call.arguments_json.as_str()}
        })).collect());
    }
    value
}

#[derive(Deserialize)]
struct WireResponse {
    choices: Vec<WireChoice>,
    usage: Option<WireUsage>,
}

#[derive(Deserialize)]
struct WireChoice {
    message: WireMessage,
}

#[derive(Deserialize)]
struct WireMessage {
    content: Option<String>,
    #[serde(default)]
    tool_calls: Vec<WireToolCall>,
}

#[derive(Deserialize)]
struct WireToolCall {
    id: String,
    #[serde(rename = "type")]
    kind: Option<String>,
    function: WireFunction,
}

#[derive(Deserialize)]
struct WireFunction {
    name: String,
    arguments: String,
}

#[derive(Deserialize)]
struct WireUsage {
    #[serde(rename = "prompt_tokens")]
    prompt: Option<u64>,
    #[serde(rename = "completion_tokens")]
    completion: Option<u64>,
    #[serde(rename = "total_tokens")]
    total: Option<u64>,
}

fn parse_completion(
    request_id: u64,
    correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
    body: &[u8],
    latency: Duration,
    attempts: usize,
) -> Result<Completion, LlmError> {
    let wire: WireResponse = serde_json::from_slice(body).map_err(|_| LlmError::MalformedJson)?;
    let choice = wire
        .choices
        .into_iter()
        .next()
        .ok_or(LlmError::UnsupportedResponse)?;
    let mut proposed_calls = BoundedVec::new();
    let mut ids = BTreeSet::new();
    for call in choice.message.tool_calls {
        if call.kind.as_deref().is_some_and(|kind| kind != "function") {
            return Err(LlmError::UnsupportedResponse);
        }
        if !ids.insert(call.id.clone()) {
            return Err(LlmError::DuplicateToolCallId);
        }
        proposed_calls
            .try_push(
                "completion.tool_calls",
                ProposedToolCall::from_model_output(
                    call.id,
                    call.function.name,
                    call.function.arguments,
                )?,
            )
            .map_err(LlmError::Boundary)?;
    }
    let mut content = BoundedVec::new();
    if let Some(text) = choice.message.content {
        content
            .try_push(
                "completion.content",
                ContentPart::Text(BoundedText::new_allow_empty("completion.text", text)?),
            )
            .map_err(LlmError::Boundary)?;
    }
    if content.is_empty() && proposed_calls.is_empty() {
        return Err(LlmError::UnsupportedResponse);
    }
    Ok(Completion {
        request_id,
        correlation_id,
        message: CompletionMessage {
            role: MessageRole::Agent,
            content,
            tool_call_id: None,
            proposed_calls,
        },
        usage: wire.usage.map(|usage| Usage {
            prompt_tokens: usage.prompt,
            completion_tokens: usage.completion,
            total_tokens: usage.total,
        }),
        latency,
        attempts,
    })
}

fn classify_reqwest(error: &reqwest::Error) -> LlmError {
    if error.is_timeout() {
        LlmError::Timeout
    } else {
        LlmError::Transport
    }
}

fn transient(error: &LlmError) -> bool {
    matches!(error, LlmError::Transport | LlmError::Timeout)
        || matches!(
            error,
            LlmError::HttpStatus(408 | 425 | 429 | 500 | 502 | 503 | 504)
        )
}

fn deterministic_backoff(request_id: u64, attempt: usize) -> Duration {
    let exponent = u32::try_from(attempt).unwrap_or(8).min(8);
    let base = 100_u64.saturating_mul(1_u64 << exponent);
    let jitter = request_id
        .wrapping_mul(6_364_136_223_846_793_005)
        .rotate_left(exponent)
        % 97;
    Duration::from_millis(base.saturating_add(jitter))
}

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

    #[tokio::test]
    async fn total_timeout_includes_waiting_for_a_concurrency_slot() {
        let limits = LlmTransportLimits {
            total_timeout: Duration::from_millis(20),
            max_concurrent_requests: 1,
            ..LlmTransportLimits::default()
        };
        let connection = AgentConfig::offline("http://127.0.0.1:9/exact", "test-key")
            .expect("test config")
            .llm;
        let client = LlmClient::new(connection, limits).expect("test client");
        let permit = Arc::clone(&client.slots)
            .acquire_owned()
            .await
            .expect("semaphore open");
        assert_eq!(
            client
                .complete(
                    &[CompletionMessage::text(MessageRole::Avatar, "hello").expect("message")],
                    &[],
                    &CancellationToken::default(),
                )
                .await
                .expect_err("slot wait must time out"),
            LlmError::Timeout
        );
        drop(permit);
    }
}