//! 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::HashSet;

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

use crate::error::{Error, Result};
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>>,
}

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

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
    }

    /// 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));

        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(e) => format!("[error] {e}"),
                };
                messages.push(Message::tool(&call.name, content));
            }
        }

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

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,
}