PicoBot/src/tools/time.rs

use async_trait::async_trait;
use chrono::{DateTime, Days, Duration, Months, Utc};
use serde_json::{Value, json};

use super::traits::{Tool, ToolResult};

pub struct TimeTool {
    default_timezone: String,
}

impl TimeTool {
    pub fn new(default_timezone: impl Into<String>) -> Self {
        Self {
            default_timezone: default_timezone.into(),
        }
    }

    fn execute_at(&self, now_utc: DateTime<Utc>, args: Value) -> ToolResult {
        match execute_time_request(now_utc, &self.default_timezone, args) {
            Ok(output) => ToolResult {
                success: true,
                output,
                error: None,
            },
            Err(error) => ToolResult {
                success: false,
                output: String::new(),
                error: Some(error),
            },
        }
    }
}

impl Default for TimeTool {
    fn default() -> Self {
        Self::new("Asia/Shanghai")
    }
}

#[async_trait]
impl Tool for TimeTool {
    fn name(&self) -> &str {
        "get_time"
    }

    fn description(&self) -> &str {
        "Get the current time or calculate a time before/after now in minutes, hours, days, months, or years. Use this for current time, relative time offsets, or when exact timestamps matter. If amount/unit/direction are omitted, it returns the current time in the default timezone. Examples: current time, 30 minutes later, 2 hours ago, 7 days later, 1 month ago, 1 year later."
    }

    fn read_only(&self) -> bool {
        true
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "direction": {
                    "type": "string",
                    "description": "Relative direction from the current time. Optional when only the current time is needed.",
                    "enum": ["past", "future", "before", "after", "ago", "later"]
                },
                "amount": {
                    "type": "integer",
                    "description": "How many units to move relative to now. Required when direction and unit are used.",
                    "minimum": 0
                },
                "unit": {
                    "type": "string",
                    "description": "Time unit for the relative offset. Required when direction and amount are used.",
                    "enum": [
                        "minute", "minutes",
                        "hour", "hours",
                        "day", "days",
                        "month", "months",
                        "year", "years"
                    ]
                },
                "timezone": {
                    "type": "string",
                    "description": "Optional IANA timezone override, for example Asia/Shanghai or Europe/Berlin. Defaults to the PicoBot configured timezone."
                }
            }
        })
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        Ok(self.execute_at(Utc::now(), args))
    }
}

fn execute_time_request(
    now_utc: DateTime<Utc>,
    default_timezone: &str,
    args: Value,
) -> Result<String, String> {
    let timezone_name = args
        .get("timezone")
        .and_then(Value::as_str)
        .unwrap_or(default_timezone);
    let timezone = timezone_name.parse::<chrono_tz::Tz>().map_err(|_| {
        format!("Invalid timezone: {timezone_name}. Expected an IANA timezone like Asia/Shanghai")
    })?;

    let now_local = now_utc.with_timezone(&timezone);
    let offset = parse_offset_request(&args)?;

    let (kind, result_local, description, offset_json) = match offset {
        None => (
            "current",
            now_local,
            format!(
                "Current time in {timezone_name} is {}",
                format_human_time(&now_local)
            ),
            Value::Null,
        ),
        Some(offset) => {
            let result_local = apply_offset(now_local, &offset)?;
            let direction_text = match offset.direction {
                OffsetDirection::Past => "before",
                OffsetDirection::Future => "after",
            };
            let unit_label = pluralized_unit(offset.unit, offset.amount);
            let description = format!(
                "{} {} {} current time in {timezone_name} is {}",
                offset.amount,
                unit_label,
                direction_text,
                format_human_time(&result_local)
            );
            (
                "offset",
                result_local,
                description,
                json!({
                    "direction": offset.direction.as_str(),
                    "amount": offset.amount,
                    "unit": offset.unit.as_str(),
                }),
            )
        }
    };

    serde_json::to_string_pretty(&json!({
        "kind": kind,
        "timezone": timezone_name,
        "current_time": now_local.to_rfc3339(),
        "result_time": result_local.to_rfc3339(),
        "unix_timestamp": result_local.timestamp(),
        "offset": offset_json,
        "description": description,
    }))
    .map_err(|error| format!("Failed to serialize time tool output: {error}"))
}

fn parse_offset_request(args: &Value) -> Result<Option<OffsetRequest>, String> {
    let direction = args.get("direction").and_then(Value::as_str);
    let amount = args.get("amount");
    let unit = args.get("unit").and_then(Value::as_str);

    if direction.is_none() && amount.is_none() && unit.is_none() {
        return Ok(None);
    }

    let direction = direction.ok_or_else(|| {
        "Missing required parameter: direction when requesting a relative time".to_string()
    })?;
    let amount = amount.and_then(Value::as_u64).ok_or_else(|| {
        "Missing required parameter: amount when requesting a relative time".to_string()
    })?;
    let amount = u32::try_from(amount)
        .map_err(|_| "amount is too large; expected a 32-bit unsigned integer".to_string())?;
    let unit = unit.ok_or_else(|| {
        "Missing required parameter: unit when requesting a relative time".to_string()
    })?;

    Ok(Some(OffsetRequest {
        direction: OffsetDirection::parse(direction)?,
        amount,
        unit: TimeUnit::parse(unit)?,
    }))
}

fn apply_offset(
    now_local: DateTime<chrono_tz::Tz>,
    offset: &OffsetRequest,
) -> Result<DateTime<chrono_tz::Tz>, String> {
    match (offset.direction, offset.unit) {
        (OffsetDirection::Future, TimeUnit::Minute) => {
            Ok(now_local + Duration::minutes(i64::from(offset.amount)))
        }
        (OffsetDirection::Past, TimeUnit::Minute) => {
            Ok(now_local - Duration::minutes(i64::from(offset.amount)))
        }
        (OffsetDirection::Future, TimeUnit::Hour) => {
            Ok(now_local + Duration::hours(i64::from(offset.amount)))
        }
        (OffsetDirection::Past, TimeUnit::Hour) => {
            Ok(now_local - Duration::hours(i64::from(offset.amount)))
        }
        (OffsetDirection::Future, TimeUnit::Day) => now_local
            .checked_add_days(Days::new(u64::from(offset.amount)))
            .ok_or_else(|| "Failed to add days to the current time".to_string()),
        (OffsetDirection::Past, TimeUnit::Day) => now_local
            .checked_sub_days(Days::new(u64::from(offset.amount)))
            .ok_or_else(|| "Failed to subtract days from the current time".to_string()),
        (OffsetDirection::Future, TimeUnit::Month) => now_local
            .checked_add_months(Months::new(offset.amount))
            .ok_or_else(|| "Failed to add months to the current time".to_string()),
        (OffsetDirection::Past, TimeUnit::Month) => now_local
            .checked_sub_months(Months::new(offset.amount))
            .ok_or_else(|| "Failed to subtract months from the current time".to_string()),
        (OffsetDirection::Future, TimeUnit::Year) => now_local
            .checked_add_months(Months::new(years_to_months(offset.amount)?))
            .ok_or_else(|| "Failed to add years to the current time".to_string()),
        (OffsetDirection::Past, TimeUnit::Year) => now_local
            .checked_sub_months(Months::new(years_to_months(offset.amount)?))
            .ok_or_else(|| "Failed to subtract years from the current time".to_string()),
    }
}

fn years_to_months(years: u32) -> Result<u32, String> {
    years
        .checked_mul(12)
        .ok_or_else(|| "amount is too large to convert years into months".to_string())
}

fn format_human_time(value: &DateTime<chrono_tz::Tz>) -> String {
    value.format("%Y-%m-%d %H:%M:%S %Z").to_string()
}

fn pluralized_unit(unit: TimeUnit, amount: u32) -> &'static str {
    match (unit, amount) {
        (TimeUnit::Minute, 1) => "minute",
        (TimeUnit::Minute, _) => "minutes",
        (TimeUnit::Hour, 1) => "hour",
        (TimeUnit::Hour, _) => "hours",
        (TimeUnit::Day, 1) => "day",
        (TimeUnit::Day, _) => "days",
        (TimeUnit::Month, 1) => "month",
        (TimeUnit::Month, _) => "months",
        (TimeUnit::Year, 1) => "year",
        (TimeUnit::Year, _) => "years",
    }
}

#[derive(Clone, Copy)]
struct OffsetRequest {
    direction: OffsetDirection,
    amount: u32,
    unit: TimeUnit,
}

#[derive(Clone, Copy)]
enum OffsetDirection {
    Past,
    Future,
}

impl OffsetDirection {
    fn parse(value: &str) -> Result<Self, String> {
        match value.trim().to_ascii_lowercase().as_str() {
            "past" | "before" | "ago" => Ok(Self::Past),
            "future" | "after" | "later" => Ok(Self::Future),
            other => Err(format!(
                "Invalid direction: {other}. Expected one of: past, future, before, after, ago, later"
            )),
        }
    }

    fn as_str(&self) -> &'static str {
        match self {
            Self::Past => "past",
            Self::Future => "future",
        }
    }
}

#[derive(Clone, Copy)]
enum TimeUnit {
    Minute,
    Hour,
    Day,
    Month,
    Year,
}

impl TimeUnit {
    fn parse(value: &str) -> Result<Self, String> {
        match value.trim().to_ascii_lowercase().as_str() {
            "minute" | "minutes" => Ok(Self::Minute),
            "hour" | "hours" => Ok(Self::Hour),
            "day" | "days" => Ok(Self::Day),
            "month" | "months" => Ok(Self::Month),
            "year" | "years" => Ok(Self::Year),
            other => Err(format!(
                "Invalid unit: {other}. Expected one of: minute, hour, day, month, year"
            )),
        }
    }

    fn as_str(&self) -> &'static str {
        match self {
            Self::Minute => "minute",
            Self::Hour => "hour",
            Self::Day => "day",
            Self::Month => "month",
            Self::Year => "year",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{TimeZone, Timelike};

    fn fixed_now() -> DateTime<Utc> {
        Utc.with_ymd_and_hms(2026, 4, 27, 4, 30, 0)
            .single()
            .unwrap()
    }

    #[test]
    fn test_get_current_time_uses_default_timezone() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(fixed_now(), json!({}));

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(payload["kind"], "current");
        assert_eq!(payload["timezone"], "Asia/Shanghai");
        assert_eq!(payload["result_time"], "2026-04-27T12:30:00+08:00");
    }

    #[test]
    fn test_get_future_hours() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(
            fixed_now(),
            json!({"direction": "future", "amount": 2, "unit": "hours"}),
        );

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(payload["result_time"], "2026-04-27T14:30:00+08:00");
        assert_eq!(payload["offset"]["direction"], "future");
        assert_eq!(payload["offset"]["unit"], "hour");
    }

    #[test]
    fn test_get_past_minutes() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(
            fixed_now(),
            json!({"direction": "ago", "amount": 15, "unit": "minute"}),
        );

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(payload["result_time"], "2026-04-27T12:15:00+08:00");
    }

    #[test]
    fn test_get_future_days_crosses_date_boundary() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(
            fixed_now(),
            json!({"direction": "after", "amount": 3, "unit": "days"}),
        );

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(payload["result_time"], "2026-04-30T12:30:00+08:00");
    }

    #[test]
    fn test_get_past_month_rolls_back_to_last_valid_day() {
        let tool = TimeTool::new("Asia/Shanghai");
        let now = Utc.with_ymd_and_hms(2026, 3, 31, 1, 0, 0).single().unwrap();
        let result = tool.execute_at(
            now,
            json!({"direction": "past", "amount": 1, "unit": "month"}),
        );

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(payload["result_time"], "2026-02-28T09:00:00+08:00");
    }

    #[test]
    fn test_get_future_year_handles_leap_day() {
        let tool = TimeTool::new("Asia/Shanghai");
        let now = Utc.with_ymd_and_hms(2024, 2, 29, 4, 0, 0).single().unwrap();
        let result = tool.execute_at(
            now,
            json!({"direction": "future", "amount": 1, "unit": "year"}),
        );

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(payload["result_time"], "2025-02-28T12:00:00+08:00");
    }

    #[test]
    fn test_timezone_override_is_used() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(fixed_now(), json!({"timezone": "Europe/Berlin"}));

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(payload["timezone"], "Europe/Berlin");
        let result_time = payload["result_time"].as_str().unwrap();
        assert!(result_time.ends_with("+02:00"));
    }

    #[test]
    fn test_missing_relative_unit_returns_error() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(fixed_now(), json!({"direction": "future", "amount": 1}));

        assert!(!result.success);
        assert_eq!(
            result.error.as_deref(),
            Some("Missing required parameter: unit when requesting a relative time")
        );
    }

    #[test]
    fn test_invalid_timezone_returns_error() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(fixed_now(), json!({"timezone": "Mars/Base"}));

        assert!(!result.success);
        assert!(result.error.unwrap().contains("Invalid timezone"));
    }

    #[test]
    fn test_zero_offset_keeps_same_local_time() {
        let tool = TimeTool::new("Asia/Shanghai");
        let result = tool.execute_at(
            fixed_now(),
            json!({"direction": "future", "amount": 0, "unit": "hours"}),
        );

        assert!(result.success);
        let payload: Value = serde_json::from_str(&result.output).unwrap();
        let result_time =
            chrono::DateTime::parse_from_rfc3339(payload["result_time"].as_str().unwrap()).unwrap();
        assert_eq!(result_time.hour(), 12);
        assert_eq!(result_time.minute(), 30);
    }
}