//! Python bindings for prebuilt agents (AgentBuilder, AgentType, GraphBasedAgent)

use pyo3::prelude::*;

use flowgentra_ai::core::agents::{
    AgentBuilder, AgentType, GraphBasedAgent, ToolSpec,
};

use crate::error::to_py_err;
use crate::get_runtime;

// ─── PyAgentType ────────────────────────────────────────────────────────────

/// Agent type: "zero_shot_react", "few_shot_react", or "conversational".
#[pyclass(name = "AgentType")]
#[derive(Clone)]
pub struct PyAgentType {
    pub(crate) inner: AgentType,
}

#[pymethods]
impl PyAgentType {
    /// Zero-shot ReAct agent (reasoning + action without examples).
    #[staticmethod]
    fn zero_shot_react() -> Self {
        PyAgentType { inner: AgentType::ZeroShotReAct }
    }

    /// Few-shot ReAct agent (with example demonstrations).
    #[staticmethod]
    fn few_shot_react() -> Self {
        PyAgentType { inner: AgentType::FewShotReAct }
    }

    /// Conversational agent (multi-turn dialogue with memory).
    #[staticmethod]
    fn conversational() -> Self {
        PyAgentType { inner: AgentType::Conversational }
    }

    fn __repr__(&self) -> String {
        format!("AgentType({})", self.inner)
    }

    fn __str__(&self) -> String {
        self.inner.to_string()
    }
}

// ─── PyToolSpec ─────────────────────────────────────────────────────────────

/// Tool specification for prebuilt agents.
///
/// Example:
///     tool = ToolSpec("search", "Search the web")
///     tool.add_parameter("query", "string")
///     tool.set_required("query")
#[pyclass(name = "ToolSpec")]
#[derive(Clone)]
pub struct PyToolSpec {
    pub(crate) inner: ToolSpec,
}

#[pymethods]
impl PyToolSpec {
    #[new]
    fn new(name: &str, description: &str) -> Self {
        PyToolSpec {
            inner: ToolSpec::new(name, description),
        }
    }

    /// Add a parameter with its type.
    fn add_parameter(&mut self, name: &str, param_type: &str) {
        self.inner.parameters.insert(name.to_string(), param_type.to_string());
    }

    /// Mark a parameter as required.
    fn set_required(&mut self, param: &str) {
        self.inner.required.push(param.to_string());
    }

    #[getter]
    fn name(&self) -> String {
        self.inner.name.clone()
    }

    #[getter]
    fn description(&self) -> String {
        self.inner.description.clone()
    }

    fn __repr__(&self) -> String {
        format!("ToolSpec(name='{}', desc='{}')", self.inner.name, self.inner.description)
    }
}

// ─── PyAgentBuilder ─────────────────────────────────────────────────────────

/// Builder for creating prebuilt agents (similar to LangChain's initialize_agent).
///
/// Example:
///     builder = AgentBuilder(AgentType.zero_shot_react())
///     builder.with_name("my_agent")
///     builder.with_llm_config("gpt-4")
///     builder.with_system_prompt("You are a helpful assistant")
///     agent = builder.build_graph()
///     result = agent.execute_input("What is 2+2?")
#[pyclass(name = "AgentBuilder")]
pub struct PyAgentBuilder {
    inner: AgentBuilder,
}

#[pymethods]
impl PyAgentBuilder {
    #[new]
    fn new(agent_type: &PyAgentType) -> Self {
        PyAgentBuilder {
            inner: AgentBuilder::new(agent_type.inner),
        }
    }

    fn with_name(&mut self, name: &str) {
        self.inner = std::mem::take(&mut self.inner).with_name(name);
    }

    fn with_llm_config(&mut self, model: &str) {
        self.inner = std::mem::take(&mut self.inner).with_llm_config(model);
    }

    fn with_temperature(&mut self, temperature: f32) {
        self.inner = std::mem::take(&mut self.inner).with_temperature(temperature);
    }

    fn with_max_tokens(&mut self, tokens: usize) {
        self.inner = std::mem::take(&mut self.inner).with_max_tokens(tokens);
    }

    fn with_tool(&mut self, tool: &PyToolSpec) {
        self.inner = std::mem::take(&mut self.inner).with_tool(tool.inner.clone());
    }

    fn with_system_prompt(&mut self, prompt: &str) {
        self.inner = std::mem::take(&mut self.inner).with_system_prompt(prompt);
    }

    fn with_memory_steps(&mut self, steps: usize) {
        self.inner = std::mem::take(&mut self.inner).with_memory_steps(steps);
    }

    fn with_evaluation(&mut self) {
        self.inner = std::mem::take(&mut self.inner).with_evaluation();
    }

    fn with_retries(&mut self, max_retries: usize) {
        self.inner = std::mem::take(&mut self.inner).with_retries(max_retries);
    }

    fn with_param(&mut self, key: &str, value: &str) {
        self.inner = std::mem::take(&mut self.inner).with_param(key, value);
    }

    /// Build the agent as a graph-based agent.
    fn build_graph(&mut self) -> PyResult<PyGraphBasedAgent> {
        let builder = std::mem::take(&mut self.inner);
        let agent = builder.build_graph().map_err(to_py_err)?;
        Ok(PyGraphBasedAgent { inner: agent })
    }

    fn __repr__(&self) -> String {
        "AgentBuilder(...)".to_string()
    }
}

// ─── PyGraphBasedAgent ──────────────────────────────────────────────────────

/// A graph-based prebuilt agent (ReAct, Conversational, etc.).
///
/// Execute with `execute_input()` for simple string in/out.
#[pyclass(name = "GraphBasedAgent")]
pub struct PyGraphBasedAgent {
    inner: GraphBasedAgent,
}

#[pymethods]
impl PyGraphBasedAgent {
    /// Execute the agent with a text input and get a text response.
    fn execute_input(&self, input: &str) -> PyResult<String> {
        let rt = get_runtime();
        let result = rt
            .block_on(self.inner.execute_input(input))
            .map_err(to_py_err)?;
        Ok(result)
    }

    /// Get the agent name.
    #[getter]
    fn name(&self) -> String {
        self.inner.config().name.clone()
    }

    /// Get node names of the underlying graph.
    fn node_names(&self) -> Vec<String> {
        self.inner.graph().node_names()
    }

    fn __repr__(&self) -> String {
        format!("GraphBasedAgent(name='{}')", self.inner.config().name)
    }
}
