pub mod types;

use std::sync::Arc;
use std::time::Instant;

use tokio::task::JoinSet;
use tokio::time;

use crate::config::SchedulerConfig;
use crate::session::{ScheduledDeliveryError, SessionManager};
use crate::storage::{
    ClaimedScheduledRun, JobRun, ScheduledOutcomeKind, ScheduledRunCompletion, ScheduledRunStatus,
    Storage,
};

pub use types::Schedule;

/// Compute the next execution time (Unix ms) for a schedule, given `from` (Unix ms).
/// Returns `None` if no next time can be determined (e.g. an invalid cron expression).
pub fn next_run_for_schedule(schedule: &Schedule, from: i64) -> Option<i64> {
    use chrono::{TimeZone, Utc};
    use std::str::FromStr;

    match schedule {
        Schedule::At { at } => Some(*at),
        Schedule::Every { every_ms } => Some(from.saturating_add(i64::try_from(*every_ms).ok()?)),
        Schedule::Cron { expr, tz } => {
            let cron_schedule = cron::Schedule::from_str(expr.as_str()).ok()?;
            let from_secs = from / 1000;
            let from_nanos = ((from % 1000) * 1_000_000) as u32;
            let from_dt = Utc.timestamp_opt(from_secs, from_nanos).single()?;
            let next_utc = if let Some(tz_str) = tz {
                let tz: chrono_tz::Tz = tz_str.parse().ok()?;
                cron_schedule
                    .after(&from_dt.with_timezone(&tz))
                    .next()?
                    .with_timezone(&Utc)
            } else {
                cron_schedule.after(&from_dt).next()?
            };
            Some(next_utc.timestamp_millis())
        }
    }
}

fn now_ms() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as i64
}

pub struct Scheduler {
    storage: Arc<Storage>,
    session_manager: Arc<SessionManager>,
    config: SchedulerConfig,
    owner: String,
    admission: crate::gateway::reload::RuntimeAdmission,
}

impl Scheduler {
    pub fn new(
        storage: Arc<Storage>,
        session_manager: Arc<SessionManager>,
        config: SchedulerConfig,
    ) -> Self {
        Self::with_admission(
            storage,
            session_manager,
            config,
            crate::gateway::reload::RuntimeAdmission::open(),
        )
    }

    pub(crate) fn with_admission(
        storage: Arc<Storage>,
        session_manager: Arc<SessionManager>,
        config: SchedulerConfig,
        admission: crate::gateway::reload::RuntimeAdmission,
    ) -> Self {
        Self {
            storage,
            session_manager,
            config,
            owner: uuid::Uuid::new_v4().to_string(),
            admission,
        }
    }

    /// Non-blocking event loop. Execution and delivery use separate bounded
    /// JoinSets so one long Agent run cannot delay other claims or outbox work.
    pub async fn run(self: Arc<Self>) {
        let poll_duration = time::Duration::from_secs(self.config.poll_interval_secs.max(1));
        let mut interval = time::interval(poll_duration);
        interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
        let max_concurrent = self.config.max_concurrent.clamp(1, 256);
        let max_delivery = max_concurrent.clamp(1, 16);
        let mut runs = JoinSet::new();
        let mut deliveries = JoinSet::new();

        tracing::info!(
            poll_interval_secs = self.config.poll_interval_secs,
            max_concurrent,
            scheduler_owner = %self.owner,
            "Scheduler started"
        );

        loop {
            tokio::select! {
                _ = interval.tick() => {}
                Some(result) = runs.join_next(), if !runs.is_empty() => {
                    if let Err(error) = result {
                        tracing::error!(error = %error, "scheduled run task panicked");
                    }
                }
                Some(result) = deliveries.join_next(), if !deliveries.is_empty() => {
                    if let Err(error) = result {
                        tracing::error!(error = %error, "scheduled delivery task panicked");
                    }
                }
            }

            while let Some(result) = runs.try_join_next() {
                if let Err(error) = result {
                    tracing::error!(error = %error, "scheduled run task panicked");
                }
            }
            while let Some(result) = deliveries.try_join_next() {
                if let Err(error) = result {
                    tracing::error!(error = %error, "scheduled delivery task panicked");
                }
            }
            if !self.admission.is_accepting() {
                continue;
            }

            let now = now_ms();
            let delivery_slots = max_delivery.saturating_sub(deliveries.len());
            if delivery_slots > 0 {
                let lease_until = now.saturating_add(180_000);
                let delivery_owner = format!("{}:delivery:{}", self.owner, uuid::Uuid::new_v4());
                match self
                    .storage
                    .claim_scheduled_deliveries(now, lease_until, &delivery_owner, delivery_slots)
                    .await
                {
                    Ok(claimed) => {
                        for run in claimed {
                            let scheduler = self.clone();
                            deliveries.spawn(async move {
                                scheduler.deliver_claimed_run(run).await;
                            });
                        }
                    }
                    Err(error) => {
                        tracing::error!(error = %error, "scheduler: failed to claim deliveries");
                    }
                }
            }

            let run_slots = max_concurrent.saturating_sub(runs.len());
            if run_slots == 0 {
                continue;
            }
            let lease_ms = self
                .config
                .execution_timeout_secs
                .saturating_add(150)
                .saturating_mul(1000)
                .min(i64::MAX as u64) as i64;
            let run_owner = format!("{}:run:{}", self.owner, uuid::Uuid::new_v4());
            match self
                .storage
                .claim_due_scheduled_runs(now, now.saturating_add(lease_ms), &run_owner, run_slots)
                .await
            {
                Ok(claimed) => {
                    for run in claimed {
                        let scheduler = self.clone();
                        runs.spawn(async move {
                            scheduler.execute_claimed_run(run).await;
                        });
                    }
                }
                Err(error) => {
                    tracing::error!(error = %error, "scheduler: failed to claim due runs");
                }
            }
        }
    }

    async fn execute_claimed_run(self: Arc<Self>, claimed: ClaimedScheduledRun) {
        let start = Instant::now();
        let job = &claimed.job;
        let (completion, agent_execution) = if let Some(_activity) = self.admission.try_enter() {
            match self.session_manager.agent_coordinator() {
                Some(coordinator) => match coordinator
                    .execute_scheduled(
                        claimed.run_id,
                        &claimed.owner,
                        &job.id,
                        &job.name,
                        job.agent_id.as_deref(),
                        &job.prompt,
                        self.config.execution_timeout_secs.max(1),
                    )
                    .await
                {
                    Ok(execution) => {
                        if execution.status == ScheduledRunStatus::Completed
                            && let Some(outcome) = execution.outcome.clone()
                        {
                            (
                                ScheduledRunCompletion {
                                    status: ScheduledRunStatus::Completed,
                                    outcome: outcome.kind,
                                    message: outcome.message,
                                    diagnostic: execution.error.clone(),
                                    duration_ms: start.elapsed().as_millis() as i64,
                                },
                                Some(execution),
                            )
                        } else {
                            let diagnostic = execution.error.clone().unwrap_or_else(|| {
                                "scheduled Agent returned ordinary text without calling complete_scheduled_run"
                                    .to_string()
                            });
                            let status = match execution.status {
                                ScheduledRunStatus::TimedOut => ScheduledRunStatus::TimedOut,
                                ScheduledRunStatus::Interrupted | ScheduledRunStatus::Cancelled => {
                                    ScheduledRunStatus::Interrupted
                                }
                                _ => ScheduledRunStatus::Failed,
                            };
                            (
                                ScheduledRunCompletion {
                                    status,
                                    outcome: ScheduledOutcomeKind::Failed,
                                    message: if status == ScheduledRunStatus::TimedOut {
                                        format!("定时任务「{}」执行超时。", job.name)
                                    } else if status == ScheduledRunStatus::Interrupted {
                                        format!("定时任务「{}」在系统关停时被中断。", job.name)
                                    } else {
                                        format!(
                                            "定时任务「{}」未可靠完成：Agent 未提交结构化运行结果。",
                                            job.name
                                        )
                                    },
                                    diagnostic: Some(diagnostic),
                                    duration_ms: start.elapsed().as_millis() as i64,
                                },
                                Some(execution),
                            )
                        }
                    }
                    Err(error) => (
                        ScheduledRunCompletion {
                            status: ScheduledRunStatus::Failed,
                            outcome: ScheduledOutcomeKind::Failed,
                            message: job.agent_id.as_deref().map_or_else(
                                || format!("定时任务「{}」未能启动 Root Agent 执行。", job.name),
                                |agent_id| {
                                    format!(
                                        "定时任务「{}」无法使用 Agent「{}」执行。请恢复该 Agent 定义，或更新任务的 agent_id。",
                                        job.name, agent_id
                                    )
                                },
                            ),
                            diagnostic: Some(error.to_string()),
                            duration_ms: start.elapsed().as_millis() as i64,
                        },
                        None,
                    ),
                },
                None => (
                    ScheduledRunCompletion {
                        status: ScheduledRunStatus::Failed,
                        outcome: ScheduledOutcomeKind::Failed,
                        message: format!("定时任务「{}」未能启动执行。", job.name),
                        diagnostic: Some("AgentCoordinator is unavailable".to_string()),
                        duration_ms: start.elapsed().as_millis() as i64,
                    },
                    None,
                ),
            }
        } else {
            (
                ScheduledRunCompletion {
                    status: ScheduledRunStatus::Interrupted,
                    outcome: ScheduledOutcomeKind::Failed,
                    message: format!("定时任务「{}」因 Gateway 重载而中断。", job.name),
                    diagnostic: Some("runtime admission closed".to_string()),
                    duration_ms: start.elapsed().as_millis() as i64,
                },
                None,
            )
        };
        let finished_at = now_ms();
        let mut commit_attempt = 0_u64;
        let commit = loop {
            let result = match agent_execution.as_ref() {
                Some(execution) => {
                    self.storage
                        .finish_scheduled_run_with_agent(
                            claimed.run_id,
                            &claimed.owner,
                            &completion,
                            &execution.agent_run_id,
                            &execution.agent_run_id,
                            execution.runtime_generation,
                            &execution.agent_terminal,
                            finished_at,
                        )
                        .await
                }
                None => {
                    self.storage
                        .finish_scheduled_run(
                            claimed.run_id,
                            &claimed.owner,
                            &completion,
                            finished_at,
                        )
                        .await
                }
            };
            match result {
                Err(error) if error.is_transient() && commit_attempt < 2 => {
                    commit_attempt += 1;
                    tracing::warn!(
                        job_id = %job.id,
                        run_id = claimed.run_id,
                        attempt = commit_attempt + 1,
                        error = %error,
                        "scheduler: retrying transient run completion commit"
                    );
                    time::sleep(time::Duration::from_millis(50 * commit_attempt)).await;
                }
                result => break result,
            }
        };
        match commit {
            Ok(true) => tracing::info!(
                job_id = %job.id,
                run_id = claimed.run_id,
                status = completion.status.as_str(),
                outcome = completion.outcome.as_str(),
                duration_ms = completion.duration_ms,
                "scheduler: run completed"
            ),
            Ok(false) => tracing::warn!(
                job_id = %job.id,
                run_id = claimed.run_id,
                "scheduler: late run result discarded"
            ),
            Err(error) => tracing::error!(
                job_id = %job.id,
                run_id = claimed.run_id,
                error = %error,
                "scheduler: failed to commit run completion"
            ),
        }
    }

    async fn deliver_claimed_run(self: Arc<Self>, run: JobRun) {
        let Some(delivery_owner) = run.delivery_lease_owner.clone() else {
            tracing::error!(
                run_id = run.id,
                "scheduler: claimed delivery has no lease owner"
            );
            return;
        };
        let result = self
            .session_manager
            .deliver_scheduled_run(&run, &delivery_owner)
            .await;
        let (delivered, permanent, error) = match result {
            Ok(()) => (true, false, None),
            Err(ScheduledDeliveryError::Transient(error)) => {
                (false, false, Some(sanitize_error(&error)))
            }
            Err(ScheduledDeliveryError::Permanent(error)) => {
                (false, true, Some(sanitize_error(&error)))
            }
        };
        if let Err(commit_error) = self
            .storage
            .complete_scheduled_delivery(
                run.id,
                &delivery_owner,
                delivered,
                permanent,
                error.as_deref(),
                now_ms(),
            )
            .await
        {
            tracing::error!(
                run_id = run.id,
                error = %commit_error,
                "scheduler: failed to commit delivery receipt"
            );
        }
    }
}

fn sanitize_error(error: &str) -> String {
    error.chars().take(1_024).collect()
}

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

    #[test]
    fn next_run_for_every_uses_claim_time() {
        assert_eq!(
            next_run_for_schedule(&Schedule::Every { every_ms: 5_000 }, 1_000),
            Some(6_000)
        );
    }

    #[test]
    fn next_run_for_at_keeps_absolute_timestamp() {
        assert_eq!(
            next_run_for_schedule(&Schedule::At { at: 2_000 }, 1_000),
            Some(2_000)
        );
    }

    #[test]
    fn cron_timezone_uses_from_argument() {
        let schedule = Schedule::Cron {
            expr: "0 0 9 * * *".to_string(),
            tz: Some("Asia/Shanghai".to_string()),
        };
        let from = chrono::DateTime::parse_from_rfc3339("2026-06-16T00:30:00Z")
            .unwrap()
            .timestamp_millis();
        let expected = chrono::DateTime::parse_from_rfc3339("2026-06-16T01:00:00Z")
            .unwrap()
            .timestamp_millis();
        assert_eq!(next_run_for_schedule(&schedule, from), Some(expected));
    }
}
