//! Provider-agnostic agent loop.
//!
//! The [`Agent`] runs the tool-use loop: send `messages + tools` to an
//! [`InferenceProvider`], dispatch any returned tool calls against a
//! [`ToolRegistry`], feed the results back, repeat until the provider returns
//! a plain assistant message or the step budget is exhausted.
//!
//! Providers (Ollama, and later OpenAI / Anthropic / OpenAI-compatible) only
//! need to implement the small [`InferenceProvider`] trait; the loop is
//! reused across all of them.

use std::collections::{HashMap, HashSet};

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::error::{Error, Result};
use crate::schema::render_slots;
use crate::tool::{SurfaceProjection, ToolRegistry, ToolSpec};

/// A single turn in the conversation. Provider-neutral.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub role: Role,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_calls: Vec<ToolCall>,
    /// When `role == Tool`, the name of the tool that produced this result.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_name: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
    Tool,
}

/// A tool invocation requested by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// Provider-supplied ID, if any. Empty string if the provider doesn't emit one.
    #[serde(default)]
    pub id: String,
    pub name: String,
    #[serde(default)]
    pub arguments: Value,
}

impl Message {
    pub fn system(text: impl Into<String>) -> Self {
        Self {
            role: Role::System,
            content: Some(text.into()),
            tool_calls: Vec::new(),
            tool_name: None,
        }
    }
    pub fn user(text: impl Into<String>) -> Self {
        Self {
            role: Role::User,
            content: Some(text.into()),
            tool_calls: Vec::new(),
            tool_name: None,
        }
    }
    pub fn tool(name: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            role: Role::Tool,
            content: Some(content.into()),
            tool_calls: Vec::new(),
            tool_name: Some(name.into()),
        }
    }
}

/// Provider abstraction: given a conversation and a tool surface, return the
/// next assistant message. Whether the message carries tool calls or a plain
/// text turn is up to the model.
#[async_trait]
pub trait InferenceProvider: Send + Sync {
    async fn chat(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Message>;
}

/// Configuration for a [`Agent`] run.
pub struct AgentConfig {
    pub max_steps: usize,
    pub projection: SurfaceProjection,
    /// Which write capabilities are granted for this session. `None` bypasses
    /// the check entirely — appropriate only for fully-trusted callers.
    pub grants: Option<HashSet<String>>,
    /// How many times one tool may fail schema validation and still be invited
    /// to try again. Past this, the diagnostic carries a stable "do not retry"
    /// suffix so a model that cannot get the arguments right stops burning the
    /// step budget on the same call.
    pub max_repairs: usize,
}

impl Default for AgentConfig {
    fn default() -> Self {
        Self {
            max_steps: 8,
            projection: SurfaceProjection::Full,
            grants: Some(HashSet::new()),
            max_repairs: 2,
        }
    }
}

pub struct Agent<P: InferenceProvider> {
    provider: P,
    config: AgentConfig,
}

impl<P: InferenceProvider> Agent<P> {
    pub fn new(provider: P) -> Self {
        Self {
            provider,
            config: AgentConfig::default(),
        }
    }

    #[must_use]
    pub fn with_max_steps(mut self, n: usize) -> Self {
        self.config.max_steps = n;
        self
    }

    #[must_use]
    pub fn with_projection(mut self, projection: SurfaceProjection) -> Self {
        self.config.projection = projection;
        self
    }

    #[must_use]
    pub fn with_grants(mut self, grants: HashSet<String>) -> Self {
        self.config.grants = Some(grants);
        self
    }

    /// Bypass capability checks. Only appropriate for fully-trusted callers.
    #[must_use]
    pub fn without_grant_checks(mut self) -> Self {
        self.config.grants = None;
        self
    }

    #[must_use]
    pub fn with_max_repairs(mut self, n: usize) -> Self {
        self.config.max_repairs = n;
        self
    }

    /// Run the loop. Returns the final assistant message plus the transcript.
    pub async fn run(
        &self,
        registry: &ToolRegistry,
        system: Option<&str>,
        user: &str,
    ) -> Result<RunOutcome> {
        let tools = registry.specs_projected(self.config.projection);

        let mut messages: Vec<Message> = Vec::new();
        if let Some(s) = system {
            messages.push(Message::system(s));
        }
        messages.push(Message::user(user));

        // Per-tool count of schema-validation failures, so a model looping on
        // the same malformed call can be told to stop rather than being
        // invited to repair forever.
        let mut repairs: HashMap<String, usize> = HashMap::new();

        for step in 0..self.config.max_steps {
            let assistant = self.provider.chat(&messages, &tools).await?;
            let has_tool_calls = !assistant.tool_calls.is_empty();
            messages.push(assistant.clone());

            if !has_tool_calls {
                return Ok(RunOutcome {
                    final_message: assistant.content.unwrap_or_default(),
                    transcript: messages,
                    steps: step + 1,
                });
            }

            for call in assistant.tool_calls {
                let result = registry
                    .call(&call.name, call.arguments, self.config.grants.as_ref())
                    .await;
                let content = match result {
                    Ok(r) => tool_result_to_text(&r),
                    Err(Error::InvalidArguments { tool, slots }) => {
                        let seen = repairs.entry(tool.clone()).or_insert(0);
                        *seen += 1;
                        self.repair_prompt(&tool, &render_slots(&slots), *seen)
                    }
                    Err(e) => format!("[error] {e}"),
                };
                messages.push(Message::tool(&call.name, content));
            }
        }

        Err(Error::Protocol(format!(
            "agent exceeded max_steps ({})",
            self.config.max_steps
        )))
    }

    /// The tool message a model sees after failing schema validation.
    ///
    /// Names the offending fields and the tool to call again, because a model
    /// handed only "invalid arguments" re-sends the same call. Past
    /// `max_repairs` it flips to a stable refusal instead, matching the "do
    /// not retry" suffix [`Error::CapabilityDenied`] uses.
    fn repair_prompt(&self, tool: &str, slots: &str, attempt: usize) -> String {
        if attempt > self.config.max_repairs {
            format!(
                "[error] invalid arguments for `{tool}`: {slots}. \
                 This call has now failed validation {attempt} times; \
                 do not retry `{tool}`."
            )
        } else {
            format!(
                "[error] invalid arguments for `{tool}`: {slots}. \
                 Correct the named fields and call `{tool}` again."
            )
        }
    }
}

fn tool_result_to_text(result: &crate::tool::ToolCallResult) -> String {
    use crate::tool::ContentPart;
    let mut out = String::new();
    for part in &result.content {
        match part {
            ContentPart::Text { text } => {
                if !out.is_empty() {
                    out.push('\n');
                }
                out.push_str(text);
            }
        }
    }
    if result.is_error {
        format!("[error] {out}")
    } else {
        out
    }
}

/// Outcome of a completed agent run.
pub struct RunOutcome {
    pub final_message: String,
    pub transcript: Vec<Message>,
    pub steps: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tool::{Tool, ToolCallResult, ToolKind};
    use serde_json::json;
    use std::sync::Mutex;

    /// Replays a fixed script of assistant turns, recording what it was told.
    struct Scripted {
        turns: Mutex<std::vec::IntoIter<Message>>,
        seen: Mutex<Vec<Message>>,
    }

    impl Scripted {
        fn new(turns: Vec<Message>) -> Self {
            Self {
                turns: Mutex::new(turns.into_iter()),
                seen: Mutex::new(Vec::new()),
            }
        }

        /// Every tool-role message the loop fed back.
        fn tool_messages(&self) -> Vec<String> {
            self.seen
                .lock()
                .unwrap()
                .iter()
                .filter(|m| m.role == Role::Tool)
                .filter_map(|m| m.content.clone())
                .collect()
        }
    }

    #[async_trait]
    impl InferenceProvider for Scripted {
        async fn chat(&self, messages: &[Message], _tools: &[ToolSpec]) -> Result<Message> {
            *self.seen.lock().unwrap() = messages.to_vec();
            self.turns
                .lock()
                .unwrap()
                .next()
                .ok_or_else(|| Error::Protocol("script exhausted".into()))
        }
    }

    fn call(name: &str, args: Value) -> Message {
        Message {
            role: Role::Assistant,
            content: None,
            tool_calls: vec![ToolCall {
                id: String::new(),
                name: name.to_string(),
                arguments: args,
            }],
            tool_name: None,
        }
    }

    fn done(text: &str) -> Message {
        Message {
            role: Role::Assistant,
            content: Some(text.to_string()),
            tool_calls: Vec::new(),
            tool_name: None,
        }
    }

    struct Greet;
    #[async_trait]
    impl Tool for Greet {
        fn name(&self) -> &'static str {
            "greet"
        }
        fn description(&self) -> &'static str {
            "greets a name"
        }
        fn kind(&self) -> ToolKind {
            ToolKind::Read
        }
        fn input_schema(&self) -> Value {
            json!({
                "type": "object",
                "properties": { "name": { "type": "string" } },
                "required": ["name"]
            })
        }
        async fn call(&self, args: Value) -> Result<ToolCallResult> {
            Ok(ToolCallResult::text(format!(
                "hello {}",
                args["name"].as_str().unwrap_or("?")
            )))
        }
    }

    fn registry() -> ToolRegistry {
        let mut r = ToolRegistry::new();
        r.register(Greet);
        r
    }

    #[tokio::test]
    async fn a_bad_call_is_repaired_on_the_next_turn() {
        let provider = Scripted::new(vec![
            call("greet", json!({})),                // forgets `name`
            call("greet", json!({ "name": "max" })), // corrects it
            done("greeted"),
        ]);
        let agent = Agent::new(provider);
        let out = agent.run(&registry(), None, "say hi").await.unwrap();

        assert_eq!(out.final_message, "greeted");
        assert_eq!(out.steps, 3);

        let fed = agent.provider.tool_messages();
        assert!(
            fed[0].contains("missing required field `name` (string)"),
            "diagnostic must name the field: {}",
            fed[0]
        );
        assert!(
            fed[0].contains("call `greet` again"),
            "diagnostic must invite the retry: {}",
            fed[0]
        );
        assert_eq!(fed[1], "hello max");
    }

    #[tokio::test]
    async fn a_model_looping_on_the_same_bad_call_is_told_to_stop() {
        let provider = Scripted::new(vec![
            call("greet", json!({})),
            call("greet", json!({})),
            call("greet", json!({})),
            done("gave up"),
        ]);
        let agent = Agent::new(provider).with_max_repairs(2);
        let out = agent.run(&registry(), None, "say hi").await.unwrap();
        assert_eq!(out.final_message, "gave up");

        let fed = agent.provider.tool_messages();
        assert!(fed[0].contains("call `greet` again"));
        assert!(fed[1].contains("call `greet` again"));
        assert!(
            fed[2].contains("do not retry `greet`"),
            "third failure must carry the stable refusal: {}",
            fed[2]
        );
    }

    #[tokio::test]
    async fn a_valid_call_never_produces_a_diagnostic() {
        let provider = Scripted::new(vec![call("greet", json!({ "name": "max" })), done("ok")]);
        let agent = Agent::new(provider);
        agent.run(&registry(), None, "say hi").await.unwrap();
        assert_eq!(agent.provider.tool_messages(), ["hello max"]);
    }
}