//! One turn: ask, run any tool calls under their limits, and return.
//! The loop is bounded every way a runaway can go: a cap on tool iterations, a repeated-call
//! detector, the two M2a limits (thinking tokens, a full context), and a tool-result cap.

use crate::baseline::messages;
use crate::config::Config;
use crate::llama::info::{CacheOutcome, cache_outcome};
use crate::llama::{ChatEvent, ChatRequest, Client, InferError};
use crate::session::{Session, SessionError};
use crate::tools::{Registry, ToolPort, cap_result, denial_text, dispatch};
use proto::{CallId, DataClass, LogRecord, Timestamp, ToolRequest, ToolResponse, TurnEvent, Usage};

/// What one turn can fail with, beyond the inference server.
#[derive(Debug)]
pub enum TurnError {
    SessionFull,
    TurnLimit,
    Infer(InferError),
    Session(SessionError),
}

impl std::fmt::Display for TurnError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TurnError::SessionFull => write!(f, "the context is full"),
            TurnError::TurnLimit => write!(f, "the turn limit was reached"),
            TurnError::Infer(e) => write!(f, "{e}"),
            TurnError::Session(e) => write!(f, "{e}"),
        }
    }
}

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

impl From<SessionError> for TurnError {
    fn from(e: SessionError) -> Self {
        TurnError::Session(e)
    }
}

/// The text to show and what the turn cost.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TurnOutcome {
    pub content: String,
    pub usage: Usage,
}

/// What a turn needs besides the session.
pub struct Runtime<'a> {
    pub cfg: &'a Config,
    pub client: &'a Client,
    pub port: &'a dyn ToolPort,
    pub registry: &'a Registry,
}

/// True for the one 400 the server sends when the prompt does not fit.
pub fn is_context_full(error: &InferError) -> bool {
    match error {
        InferError::Http { status: 400, body } => {
            let Ok(value) = serde_json::from_str::<serde_json::Value>(body) else {
                return false;
            };
            matches!(
                value
                    .get("error")
                    .and_then(|e| e.get("type"))
                    .and_then(|t| t.as_str()),
                Some("exceed_context_size_error")
            )
        }
        _ => false,
    }
}

fn map_event(event: &ChatEvent, on_event: &mut dyn FnMut(&TurnEvent)) {
    match event {
        ChatEvent::Queued { ahead } => on_event(&TurnEvent::Queued {
            ahead: u64::try_from(*ahead).unwrap_or(u64::MAX),
        }),
        ChatEvent::Waiting { slot_busy } => on_event(&TurnEvent::Waiting {
            slot_busy: *slot_busy,
        }),
        ChatEvent::Progress {
            total,
            cache,
            processed,
        } => on_event(&TurnEvent::Progress {
            total: *total,
            cache: *cache,
            processed: *processed,
        }),
        ChatEvent::Reasoning(text) => on_event(&TurnEvent::Reasoning { text: text.clone() }),
        ChatEvent::Content(text) => on_event(&TurnEvent::Content { text: text.clone() }),
        ChatEvent::ThinkingCapped { tokens } => {
            on_event(&TurnEvent::ThinkingCapped { tokens: *tokens })
        }
        ChatEvent::Retrying {
            attempt,
            after_ms,
            error,
        } => on_event(&TurnEvent::Retrying {
            attempt: *attempt,
            after_ms: *after_ms,
            error: error.clone(),
        }),
        ChatEvent::ToolCallDelta { .. } => {}
    }
}

pub fn run_turn(
    session: &mut Session,
    rt: &Runtime<'_>,
    content: &str,
    on_event: &mut dyn FnMut(&TurnEvent),
) -> Result<TurnOutcome, TurnError> {
    session
        .append(LogRecord::User {
            time: Timestamp::now(),
            content: content.to_string(),
        })
        .map_err(TurnError::Session)?;

    let mut seen: Vec<(String, String)> = Vec::new();
    let mut iterations: u32 = 0;

    loop {
        let baseline = session.baseline();
        let request = ChatRequest {
            slot: rt.cfg.slots.main,
            messages: messages(baseline, session.records()),
            tools: baseline.tools.clone(),
            thinking: true,
        };
        let previous = session.last_usage();

        let completion = match rt
            .client
            .chat_with_retry(&request, &mut |event| map_event(event, on_event))
        {
            Ok(completion) => completion,
            Err(e) => {
                if is_context_full(&e) {
                    return Err(TurnError::SessionFull);
                }
                return Err(TurnError::Infer(e));
            }
        };

        session
            .append(LogRecord::Assistant {
                time: Timestamp::now(),
                content: completion.content.clone(),
                reasoning_content: completion.reasoning_content.clone(),
                tool_calls: completion.tool_calls.clone(),
            })
            .map_err(TurnError::Session)?;

        let usage = Usage {
            cache_n: completion.timings.cache_n,
            prompt_n: completion.timings.prompt_n,
            predicted_n: completion.timings.predicted_n,
            reasoning_tokens: completion.reasoning_tokens,
            thinking_capped: completion.thinking_capped,
        };
        session
            .append(LogRecord::Usage {
                time: Timestamp::now(),
                cache_n: usage.cache_n,
                prompt_n: usage.prompt_n,
                predicted_n: usage.predicted_n,
                reasoning_tokens: usage.reasoning_tokens,
                thinking_capped: usage.thinking_capped,
            })
            .map_err(TurnError::Session)?;

        // The recordings were made in separate conversations, so their cache numbers do not line up.
        if let Some(previous) = previous {
            let previous = crate::llama::Timings {
                cache_n: previous.cache_n,
                prompt_n: previous.prompt_n,
                predicted_n: previous.predicted_n,
            };
            if let CacheOutcome::Loss { expected, got } =
                cache_outcome(&previous, &completion.timings)
            {
                session
                    .append(LogRecord::CacheLoss {
                        time: Timestamp::now(),
                        expected,
                        got,
                    })
                    .map_err(TurnError::Session)?;
                on_event(&TurnEvent::CacheLoss { expected, got });
            }
        }

        if completion.tool_calls.is_empty() {
            return Ok(TurnOutcome {
                content: completion.content.clone().unwrap_or_default(),
                usage,
            });
        }

        iterations += 1;
        if iterations > rt.cfg.r#loop.tool_iterations {
            return Err(TurnError::TurnLimit);
        }

        for call in &completion.tool_calls {
            on_event(&TurnEvent::ToolCallStarted {
                name: call.name.clone(),
            });
            let call_id = session.next_call();

            // A repeated identical (name, arguments) pair is not run the first time; a second repeat stops the turn.
            let key = (call.name.clone(), call.arguments.clone());
            let times = seen.iter().filter(|pair| **pair == key).count();
            if rt.cfg.r#loop.repeat_detection && times >= 2 {
                return Err(TurnError::TurnLimit);
            }
            seen.push(key);

            let (text, class, untrusted) = if rt.cfg.r#loop.repeat_detection && times == 1 {
                (
                    format!(
                        "The `{}` tool call was already called with these arguments this turn.",
                        call.name
                    ),
                    DataClass::Public,
                    false,
                )
            } else {
                run_call(rt, session, call, call_id, on_event)
            };

            let (text, truncated) = cap_result(&text, rt.cfg.r#loop.tool_result_cap);
            session
                .append(LogRecord::ToolResult {
                    time: Timestamp::now(),
                    call: call_id,
                    tool_call_id: call.id.clone(),
                    content: text,
                    class,
                    untrusted,
                    truncated,
                })
                .map_err(TurnError::Session)?;
            on_event(&TurnEvent::ToolResult {
                name: call.name.clone(),
                class,
                truncated,
            });
        }
    }
}

/// Resolves one tool call to the text, class and trust flag of its result.
fn run_call(
    rt: &Runtime<'_>,
    session: &Session,
    call: &proto::ToolCall,
    call_id: CallId,
    on_event: &mut dyn FnMut(&TurnEvent),
) -> (String, DataClass, bool) {
    match dispatch(rt.registry, &call.name, &call.arguments) {
        crate::tools::Dispatch::Local(text) => (text, DataClass::Public, false),
        crate::tools::Dispatch::Port { tool, arguments } => {
            let request = ToolRequest {
                session: session.id().clone(),
                call: call_id,
                tool,
                arguments,
            };
            let response = rt.port.call(&request, &mut |pending| {
                on_event(&TurnEvent::ApprovalPending {
                    approval: pending.approval,
                    tool: request.tool.clone(),
                    expires: pending.expires,
                });
            });
            match response {
                ToolResponse::Result {
                    content,
                    class,
                    untrusted,
                    ..
                } => (content, class, untrusted),
                ToolResponse::Failed { message } => (
                    format!("The tool failed: {message}"),
                    DataClass::Public,
                    false,
                ),
                ToolResponse::Denied { reason } => {
                    on_event(&TurnEvent::ToolDenied {
                        name: request.tool.clone(),
                        reason,
                    });
                    (denial_text(reason).to_string(), DataClass::Public, false)
                }
                ToolResponse::PendingApproval { .. } => (
                    "The tool failed: the tool broker gave no final answer".to_string(),
                    DataClass::Public,
                    false,
                ),
            }
        }
    }
}