use crate::agent::Agent;
use crate::entities::task::Entity as Task;
use crate::entities::task_run;
use dashmap::DashMap;
use sea_orm::{DatabaseConnection, EntityTrait, Set};
use tokio_cron_scheduler::{Job, JobScheduler};
use uuid::Uuid;

pub struct Scheduler {
    scheduler: JobScheduler,
    db: DatabaseConnection,
    tasks_to_jobs: DashMap<Uuid, Uuid>,
    zen_api_key: Option<String>,
    tavily_api_key: Option<String>,
}

impl Scheduler {
    pub async fn new(
        db: DatabaseConnection,
        zen_api_key: Option<String>,
        tavily_api_key: Option<String>,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let scheduler = JobScheduler::new().await?;
        scheduler.start().await?;
        Ok(Self {
            scheduler,
            db,
            tasks_to_jobs: DashMap::new(),
            zen_api_key,
            tavily_api_key,
        })
    }

    pub async fn add_task_job(
        &self,
        task_id: Uuid,
        cron_expr: &str,
    ) -> Result<(), Box<dyn std::error::Error>> {
        // Remove existing job if any
        if let Some((_, old_job_id)) = self.tasks_to_jobs.remove(&task_id) {
            let _ = self.scheduler.remove(&old_job_id).await;
        }

        let db = self.db.clone();
        let zen_key = self.zen_api_key.clone();
        let tavily_key = self.tavily_api_key.clone();

        let job = Job::new_async(cron_expr, move |_uuid, _l| {
            let db = db.clone();
            let zen_key = zen_key.clone();
            let tavily_key = tavily_key.clone();
            Box::pin(async move {
                if let Err(e) = Self::run_task(db, zen_key, tavily_key, task_id).await {
                    eprintln!("Error in scheduled task {}: {}", task_id, e);
                }
            })
        })?;

        let job_id = self.scheduler.add(job).await?;
        self.tasks_to_jobs.insert(task_id, job_id);

        Ok(())
    }

    pub async fn remove_task_job(&self, task_id: Uuid) -> Result<(), Box<dyn std::error::Error>> {
        if let Some((_, job_id)) = self.tasks_to_jobs.remove(&task_id) {
            self.scheduler.remove(&job_id).await?;
        }
        Ok(())
    }

    async fn run_task(
        db: DatabaseConnection,
        zen_key: Option<String>,
        tavily_key: Option<String>,
        task_id: Uuid,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let task = Task::find_by_id(task_id)
            .one(&db)
            .await?
            .ok_or("Task not found")?;

        // Create a new run entry
        let run_id = Uuid::new_v4();
        let run = task_run::ActiveModel {
            id: Set(run_id),
            task_id: Set(task_id),
            status: Set("running".to_string()),
            logs: Set(String::new()),
            answer: Set(None),
            created_at: Set(chrono::Utc::now().into()),
        };

        use sea_orm::ActiveModelTrait;
        run.insert(&db).await?;

        // Start agent in background
        let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?;

        tokio::spawn(async move {
            let (logs, answer) = match agent.run().await {
                Ok(res) => res,
                Err(e) => (format!("Scheduled run failed: {}", e), None),
            };

            let run_complete = task_run::ActiveModel {
                id: Set(run_id),
                status: Set("completed".to_string()),
                logs: Set(logs),
                answer: Set(answer),
                ..Default::default()
            };
            let _ = run_complete.update(&db).await;
        });

        Ok(())
    }
}