potato-type 0.26.0

pub mod error;

pub use crate::error::TypeError;
use pyo3::prelude::*;
use schemars::JsonSchema;
use serde::de::Error;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::any::type_name;
use std::fmt;
use std::fmt::Display;
use std::path::{Path, PathBuf};
use tracing::error;
pub mod anthropic;
pub mod common;
pub mod google;
pub mod openai;
pub mod prompt;
pub mod spec;
pub mod tools;
pub mod traits;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[pyclass(from_py_object)]
pub enum Model {
    Undefined,
}

impl Model {
    pub fn as_str(&self) -> &str {
        match self {
            Model::Undefined => "undefined",
        }
    }

    pub fn from_string(s: &str) -> Result<Self, TypeError> {
        match s.to_lowercase().as_str() {
            "undefined" => Ok(Model::Undefined),
            _ => Err(TypeError::UnknownModelError(s.to_string())),
        }
    }
}

pub enum Common {
    Undefined,
}

impl Common {
    pub fn as_str(&self) -> &str {
        match self {
            Common::Undefined => "undefined",
        }
    }
}

impl Display for Common {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[pyclass(from_py_object, eq, eq_int)]
pub enum Provider {
    OpenAI,
    Gemini,
    Google,
    Vertex,
    Anthropic,
    GoogleAdk,
    Undefined, // Added Undefined for better error handling
}

impl Provider {
    pub fn from_string(s: &str) -> Result<Self, TypeError> {
        match s.to_lowercase().as_str() {
            "openai" => Ok(Provider::OpenAI),
            "gemini" => Ok(Provider::Gemini),
            "google" => Ok(Provider::Google),
            "vertex" => Ok(Provider::Vertex),
            "anthropic" => Ok(Provider::Anthropic),
            "google_adk" => Ok(Provider::GoogleAdk),
            "undefined" => Ok(Provider::Undefined), // Handle undefined case
            _ => Err(TypeError::UnknownProviderError(s.to_string())),
        }
    }

    pub const DEFAULT_ENV_VAR: &'static str = "POTATO_HEAD_DEFAULT_PROVIDER";

    /// Reads `POTATO_HEAD_DEFAULT_PROVIDER` from the environment.
    ///
    /// - Unset, empty, or whitespace-only -> `Ok(None)`
    /// - Set and parseable -> `Ok(Some(provider))`
    /// - Set but unparseable -> `Err(TypeError::UnknownProviderError(...))`
    pub fn from_env_default() -> Result<Option<Provider>, TypeError> {
        match std::env::var(Self::DEFAULT_ENV_VAR) {
            Ok(s) => {
                let trimmed = s.trim();
                if trimmed.is_empty() {
                    Ok(None)
                } else {
                    Provider::from_string(trimmed).map(Some)
                }
            }
            Err(_) => Ok(None),
        }
    }

    /// Resolve a provider from an optional explicit string, falling back to the env var,
    /// then erroring with `MissingProviderError` if neither is available.
    pub fn resolve(provided: Option<&str>) -> Result<Provider, TypeError> {
        if let Some(s) = provided.map(str::trim).filter(|s| !s.is_empty()) {
            return Provider::from_string(s);
        }

        Self::from_env_default()?.ok_or(TypeError::MissingProviderError(Self::DEFAULT_ENV_VAR))
    }

    /// PyO3-friendly resolution: accepts an optional `Bound<'_, PyAny>` (Provider enum or string).
    /// `None` or `py.None()` triggers env-var fallback.
    pub fn resolve_from_py(provider: Option<&Bound<'_, PyAny>>) -> Result<Provider, TypeError> {
        match provider {
            Some(p) if !p.is_none() => Self::extract_provider(p),
            _ => Self::from_env_default()?
                .ok_or(TypeError::MissingProviderError(Self::DEFAULT_ENV_VAR)),
        }
    }

    /// Extract provider from a PyAny object
    ///
    /// # Arguments
    /// * `provider` - PyAny object
    ///
    /// # Returns
    /// * `Result<Provider, AgentError>` - Result
    ///
    /// # Errors
    /// * `AgentError` - Error
    pub fn extract_provider(provider: &Bound<'_, PyAny>) -> Result<Provider, TypeError> {
        match provider.is_instance_of::<Provider>() {
            true => Ok(provider.extract::<Provider>().inspect_err(|e| {
                error!("Failed to extract provider: {}", e);
            })?),
            false => {
                let provider = provider.extract::<String>().unwrap();
                Ok(Provider::from_string(&provider).inspect_err(|e| {
                    error!("Failed to convert string to provider: {}", e);
                })?)
            }
        }
    }

    pub fn as_str(&self) -> &str {
        match self {
            Provider::OpenAI => "openai",
            Provider::Gemini => "gemini",
            Provider::Vertex => "vertex",
            Provider::Google => "google",
            Provider::Anthropic => "anthropic",
            Provider::GoogleAdk => "google_adk",
            Provider::Undefined => "undefined", // Added Undefined case
        }
    }
}

impl Display for Provider {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

#[pyclass(from_py_object, eq, eq_int)]
#[derive(Debug, PartialEq, Clone)]
pub enum SaveName {
    Prompt,
}

#[pymethods]
impl SaveName {
    #[staticmethod]
    pub fn from_string(s: &str) -> Option<Self> {
        match s {
            "prompt" => Some(SaveName::Prompt),

            _ => None,
        }
    }

    pub fn as_string(&self) -> &str {
        match self {
            SaveName::Prompt => "prompt",
        }
    }

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

impl Display for SaveName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.as_string())
    }
}

impl AsRef<Path> for SaveName {
    fn as_ref(&self) -> &Path {
        match self {
            SaveName::Prompt => Path::new("prompt"),
        }
    }
}

// impl PathBuf: From<SaveName>
impl From<SaveName> for PathBuf {
    fn from(save_name: SaveName) -> Self {
        PathBuf::from(save_name.as_ref())
    }
}

/// A trait for structured output types that can be used with potatohead prompts agents and workflows.
///
/// # Example
/// ```rust
/// use potatohead_macro::StructureOutput;
/// use serde::{Serialize, Deserialize};
/// use schemars::JsonSchema;
///
/// #[derive(Serialize, Deserialize, JsonSchema)]
/// struct MyOutput {
///     message: String,
///     value: i32,
/// }
///
/// impl StructuredOutput for MyOutput {}
///
/// let schema = MyOutput::get_structured_output_schema();
/// ```
pub trait StructuredOutput: for<'de> serde::Deserialize<'de> + JsonSchema {
    fn type_name() -> &'static str {
        type_name::<Self>().rsplit("::").next().unwrap_or("Unknown")
    }

    /// Validates and deserializes a JSON value into its struct type.
    ///
    /// # Arguments
    /// * `value` - The JSON value to deserialize
    ///
    /// # Returns
    /// * `Result<Self, serde_json::Error>` - The deserialized value or error
    fn model_validate_json_value(value: &Value) -> Result<Self, serde_json::Error> {
        match &value {
            Value::String(json_str) => Self::model_validate_json_str(json_str),
            Value::Object(_) => {
                // Content is already a JSON object
                serde_json::from_value(value.clone())
            }
            _ => {
                // If the value is not a string or object, we cannot deserialize it
                Err(Error::custom("Expected a JSON string or object"))
            }
        }
    }

    fn model_validate_json_str(value: &str) -> Result<Self, serde_json::Error> {
        serde_json::from_str(value)
    }

    /// Generates an OpenAI-compatible JSON schema.
    ///
    /// # Returns
    /// * `Value` - The JSON schema wrapped in OpenAI's format
    fn get_structured_output_schema() -> Value {
        let schema = ::schemars::schema_for!(Self);
        schema.into()
    }
    // add fallback parsing logic
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[pyclass(from_py_object)]
pub enum SettingsType {
    GoogleChat,
    OpenAIChat,
    ModelSettings,
    Anthropic,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;

    static ENV_LOCK: Mutex<()> = Mutex::new(());

    fn with_env_var<F: FnOnce()>(value: Option<&str>, f: F) {
        let _guard = ENV_LOCK.lock().unwrap();
        let prev = std::env::var(Provider::DEFAULT_ENV_VAR).ok();
        match value {
            Some(v) => std::env::set_var(Provider::DEFAULT_ENV_VAR, v),
            None => std::env::remove_var(Provider::DEFAULT_ENV_VAR),
        }
        f();
        match prev {
            Some(v) => std::env::set_var(Provider::DEFAULT_ENV_VAR, v),
            None => std::env::remove_var(Provider::DEFAULT_ENV_VAR),
        }
    }

    #[test]
    fn test_provider_google_adk_round_trip() {
        let p = Provider::from_string("google_adk").unwrap();
        assert_eq!(p, Provider::GoogleAdk);
        assert_eq!(p.as_str(), "google_adk");
    }

    #[test]
    fn test_provider_all_variants_round_trip() {
        for (s, variant) in [
            ("openai", Provider::OpenAI),
            ("gemini", Provider::Gemini),
            ("google", Provider::Google),
            ("vertex", Provider::Vertex),
            ("anthropic", Provider::Anthropic),
            ("google_adk", Provider::GoogleAdk),
            ("undefined", Provider::Undefined),
        ] {
            let parsed = Provider::from_string(s).unwrap();
            assert_eq!(parsed, variant);
            assert_eq!(parsed.as_str(), s);
        }
    }

    #[test]
    fn test_provider_unknown_string_errors() {
        assert!(Provider::from_string("not_a_provider").is_err());
    }

    #[test]
    fn resolve_explicit_wins_over_env() {
        with_env_var(Some("anthropic"), || {
            let p = Provider::resolve(Some("openai")).unwrap();
            assert_eq!(p, Provider::OpenAI);
        });
    }

    #[test]
    fn resolve_uses_env_when_explicit_missing() {
        with_env_var(Some("openai"), || {
            let p = Provider::resolve(None).unwrap();
            assert_eq!(p, Provider::OpenAI);
        });
    }

    #[test]
    fn resolve_uses_env_when_explicit_empty_string() {
        with_env_var(Some("openai"), || {
            let p = Provider::resolve(Some("   ")).unwrap();
            assert_eq!(p, Provider::OpenAI);
        });
    }

    #[test]
    fn resolve_empty_env_treated_as_unset() {
        with_env_var(Some("   "), || {
            let err = Provider::resolve(None).unwrap_err();
            assert!(matches!(err, TypeError::MissingProviderError(_)));
        });
    }

    #[test]
    fn resolve_invalid_env_is_hard_error() {
        with_env_var(Some("not_a_provider"), || {
            let err = Provider::resolve(None).unwrap_err();
            assert!(matches!(err, TypeError::UnknownProviderError(_)));
        });
    }

    #[test]
    fn resolve_no_explicit_no_env_returns_missing() {
        with_env_var(None, || {
            let err = Provider::resolve(None).unwrap_err();
            match err {
                TypeError::MissingProviderError(name) => {
                    assert_eq!(name, "POTATO_HEAD_DEFAULT_PROVIDER");
                }
                other => panic!("expected MissingProviderError, got {:?}", other),
            }
        });
    }

    #[test]
    fn from_env_default_unset_returns_none() {
        with_env_var(None, || {
            assert!(Provider::from_env_default().unwrap().is_none());
        });
    }

    #[test]
    fn from_env_default_set_and_parseable() {
        with_env_var(Some("gemini"), || {
            assert_eq!(
                Provider::from_env_default().unwrap(),
                Some(Provider::Gemini)
            );
        });
    }

    #[test]
    fn from_env_default_set_but_invalid_errors() {
        with_env_var(Some("garbage"), || {
            assert!(Provider::from_env_default().is_err());
        });
    }
}