//! Tool-calling loop driving the model against a [`Sandbox`].
//!
//! The loop sends the conversation and the available tools to OpenRouter,
//! executes any requested tool call inside the sandbox and feeds the results
//! back, until the model produces a final message or the iteration budget is
//! exhausted.

use std::str::FromStr;
use std::time::Instant;

use anyhow::Context;
use serde_json::Value;
use tracing::{info, warn};

use crate::{
    open_router::{Message, OpenRouterClient, Role, Tool, ToolCall, ToolChoice},
    sandbox::{Sandbox, tools},
    text::{excerpt, truncate_bytes},
};

/// Number of characters of the arguments of a tool call kept in the logs.
const MAX_LOGGED_ARGUMENTS: usize = 200;

/// Number of characters of a tool result kept in the logs.
///
/// Sizes alone make a broken tool invisible: a command that fails returns a short
/// `error: …` instead of the expected content, which is exactly what this excerpt
/// makes obvious.
const MAX_LOGGED_OUTPUT: usize = 120;

/// Maximum number of bytes of a tool result handed to the model.
///
/// A result stays in the conversation and is re-sent on every following
/// iteration, so an unbounded one (a whole file, a match on every line) inflates
/// the context and the cost of the whole rest of the run — and can overflow the
/// model context window outright.
const MAX_TOOL_OUTPUT_BYTES: usize = 32 * 1024;

/// Asked of the model when its answer does not satisfy the caller's contract.
const RETRY_PROMPT: &str = "
    Your answer is not valid: {error}

    Answer again with the requested format only: a raw json object, without any
    markdown code fence and without any text before or after it.
";

/// Asked of the model on the last turn, in place of one more exploration turn.
///
/// Without it, a model still exploring on its last turn would ask for another
/// tool, and the run would end on the iteration budget with no answer at all.
const FINAL_PROMPT: &str = "
    You are out of turns: this is your last one. Answer now, with the requested
    format, from what you have already gathered, without calling any tool.
";

/// Final output of an agent run.
pub struct AgentResult<A> {
    /// Answer of the model, parsed into the type the caller asked for.
    pub answer: A,
    pub cost: Option<f64>,
    pub iterations: usize,
}

/// Runs the tool-calling loop until the model answers or `max_iterations` is
/// reached.
///
/// `tool_definitions` is the set of tools the model may call; it is selected by
/// the caller based on the webhook action (see [`tools::for_webhook`]).
///
/// The answer of the model is parsed into `A`, which is how the caller states the
/// format it expects: while [`FromStr`] rejects the answer, the error is sent back
/// to the model so that it can correct itself, which costs an iteration. The
/// parsing happens here because that is where the conversation and the remaining
/// budget are at hand.
pub async fn run<A>(
    open_router: &OpenRouterClient,
    sandbox: &Sandbox,
    tool_definitions: Vec<Tool>,
    system_prompt: &str,
    user_prompt: &str,
    max_iterations: usize,
) -> anyhow::Result<AgentResult<A>>
where
    A: FromStr,
    A::Err: std::fmt::Display,
{
    let mut messages = vec![
        Message::new(Role::System, system_prompt),
        Message::new(Role::User, user_prompt),
    ];

    info!(
        tools = %tool_definitions
            .iter()
            .map(|tool| tool.function.name.as_str())
            .collect::<Vec<_>>()
            .join(", "),
        max_iterations,
        "Starting tool-calling loop"
    );

    let mut total_cost = 0.0_f64;
    let mut has_cost = false;
    let mut last_rejection: Option<String> = None;

    for iteration in 1..=max_iterations {
        // Le dernier tour n'est plus un tour d'exploration : le modèle doit rendre
        // sa réponse, avec ce qu'il a déjà vu.
        let last = iteration == max_iterations;

        if last {
            info!(iteration, "Last turn: asking for the final answer");
            messages.push(Message::new(Role::User, FINAL_PROMPT));
        }

        let started = Instant::now();
        let response = open_router
            .chat_with_tools(
                &messages,
                &tool_definitions,
                if last {
                    ToolChoice::None
                } else {
                    ToolChoice::Auto
                },
            )
            .await?;

        if let Some(cost) = response.cost {
            total_cost += cost;
            has_cost = true;
        }

        info!(
            iteration,
            max_iterations,
            last_turn = last,
            tool_calls = response.tool_calls.len(),
            answer_bytes = response.message.as_deref().map_or(0, str::len),
            elapsed_ms = started.elapsed().as_millis() as u64,
            "Model answered"
        );

        // Au dernier tour, les outils sont refusés : même si le fournisseur les
        // renvoie malgré tout, on ne les exécute pas et on tente la réponse.
        if last || response.tool_calls.is_empty() {
            let answer = response.message.unwrap_or_default();

            match answer.parse::<A>() {
                Ok(parsed) => {
                    return Ok(AgentResult {
                        answer: parsed,
                        cost: has_cost.then_some(total_cost),
                        iterations: iteration,
                    });
                }
                Err(err) => {
                    warn!(iteration, %err, "Model answer was rejected, asking for another one");

                    // Keep the rejected answer in the history, so that the model
                    // sees what it has to fix.
                    messages.push(Message::new(Role::Assistant, answer));
                    messages.push(Message::new(
                        Role::User,
                        RETRY_PROMPT.replace("{error}", &err.to_string()),
                    ));

                    last_rejection = Some(err.to_string());
                    continue;
                }
            }
        }

        messages.push(Message::assistant_with_tool_calls(
            response.message.unwrap_or_default(),
            response.tool_calls.clone(),
        ));

        for call in &response.tool_calls {
            let started = Instant::now();
            let (content, truncated) = execute(sandbox, call).await;

            info!(
                iteration,
                tool = call.name(),
                arguments = %excerpt(call.arguments_json(), MAX_LOGGED_ARGUMENTS),
                output_bytes = content.len(),
                truncated,
                // Les sauts de ligne casseraient la lisibilité d'une ligne de log.
                output = %excerpt(&content, MAX_LOGGED_OUTPUT).replace('\n', " "),
                elapsed_ms = started.elapsed().as_millis() as u64,
                "Tool call finished"
            );

            messages.push(Message::tool_response(call.id(), content));
        }
    }

    warn!(
        max_iterations,
        "Agent used its last turn without a valid answer"
    );

    match last_rejection {
        Some(error) => anyhow::bail!(
            "agent exceeded the maximum number of iterations ({max_iterations}), \
             last answer rejected: {error}"
        ),
        None => anyhow::bail!("agent exceeded the maximum number of iterations ({max_iterations})"),
    }
}

/// Runs a tool call and returns its result for the model, along with whether the
/// result had to be truncated.
async fn execute(sandbox: &Sandbox, call: &ToolCall) -> (String, bool) {
    let args = match parse_args(call) {
        Ok(args) => args,
        Err(err) => return cap_output(format!("error: {err}")),
    };

    match tools::dispatch(sandbox, call.name(), &args).await {
        Ok(output) => cap_output(output),
        Err(err) => cap_output(format!("error: {err}")),
    }
}

/// Bounds [`MAX_TOOL_OUTPUT_BYTES`] of a tool result, telling the model what was
/// dropped so that it can narrow its request.
fn cap_output(mut output: String) -> (String, bool) {
    let total = output.len();

    if !truncate_bytes(&mut output, MAX_TOOL_OUTPUT_BYTES) {
        return (output, false);
    }

    let kept = output.len();
    output.push_str(&format!(
        "\n… output truncated: {total} bytes in total, the first {kept} are shown. \
         Narrow the request (path, pattern or line range) to see the rest."
    ));

    (output, true)
}

fn parse_args(call: &ToolCall) -> anyhow::Result<Value> {
    let raw = call.arguments_json().trim();
    if raw.is_empty() {
        return Ok(Value::Object(serde_json::Map::new()));
    }

    serde_json::from_str(raw)
        .with_context(|| format!("invalid arguments for tool `{}`", call.name()))
}

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

    /// Builds a tool call the way the API returns one, so the fixture also
    /// covers deserialization.
    fn tool_call(name: &str, arguments: &str) -> ToolCall {
        serde_json::from_value(json!({
            "id": "call_1",
            "type": "function",
            "function": { "name": name, "arguments": arguments }
        }))
        .unwrap()
    }

    #[test]
    fn parse_args_accepts_empty_arguments() {
        let call = tool_call("ls", "");
        assert_eq!(
            parse_args(&call).unwrap(),
            Value::Object(serde_json::Map::new())
        );
    }

    #[test]
    fn parse_args_parses_json_object() {
        let call = tool_call("ls", r#"{"path":"src"}"#);
        assert_eq!(parse_args(&call).unwrap()["path"], "src");
    }

    #[test]
    fn parse_args_rejects_invalid_json() {
        let call = tool_call("ls", "not json");
        assert!(parse_args(&call).is_err());
    }

    #[test]
    fn a_small_tool_output_is_kept_as_is() {
        let (content, truncated) = cap_output(String::from("src/main.rs"));

        assert_eq!(content, "src/main.rs");
        assert!(!truncated);
    }

    #[test]
    fn a_large_tool_output_is_truncated_and_told_to_the_model() {
        let (content, truncated) = cap_output("a".repeat(MAX_TOOL_OUTPUT_BYTES + 1));

        assert!(truncated);
        assert!(content.starts_with(&"a".repeat(MAX_TOOL_OUTPUT_BYTES)));
        assert!(content.contains("output truncated"));
    }

    #[test]
    fn a_truncated_tool_output_stays_valid_utf8() {
        let (content, truncated) = cap_output("é".repeat(MAX_TOOL_OUTPUT_BYTES));

        assert!(truncated);
        assert!(content.starts_with('é'));
    }
}