use std::collections::HashMap;
use std::sync::Mutex;

use chrono::{NaiveDateTime, Utc};
use serde::Serialize;

#[derive(Debug, Clone, Serialize)]
pub struct TaskProgress {
    pub current: u64,
    pub total: u64,
    pub message: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct TaskInfo {
    pub id: String,
    pub task_type: String,
    pub status: TaskStatus,
    pub progress: Option<TaskProgress>,
    pub started_at: NaiveDateTime,
    pub completed_at: Option<NaiveDateTime>,
    pub result: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum TaskStatus {
    Pending,
    Running,
    Completed,
    Failed,
}

pub struct TaskManager {
    tasks: Mutex<HashMap<String, TaskInfo>>,
}

impl Default for TaskManager {
    fn default() -> Self {
        Self::new()
    }
}

impl TaskManager {
    pub fn new() -> Self {
        Self {
            tasks: Mutex::new(HashMap::new()),
        }
    }

    /// Register a new task as Running. Returns the task ID.
    pub fn register(&self, task_type: &str) -> String {
        let id = uuid::Uuid::new_v4().to_string();
        let info = TaskInfo {
            id: id.clone(),
            task_type: task_type.to_string(),
            status: TaskStatus::Running,
            progress: None,
            started_at: Utc::now().naive_utc(),
            completed_at: None,
            result: None,
        };
        self.tasks.lock().unwrap().insert(id.clone(), info);
        id
    }

    /// Register a new task as Pending (queued, not yet running). Returns the task ID.
    pub fn register_pending(&self, task_type: &str) -> String {
        let id = uuid::Uuid::new_v4().to_string();
        let info = TaskInfo {
            id: id.clone(),
            task_type: task_type.to_string(),
            status: TaskStatus::Pending,
            progress: None,
            started_at: Utc::now().naive_utc(),
            completed_at: None,
            result: None,
        };
        self.tasks.lock().unwrap().insert(id.clone(), info);
        id
    }

    /// Transition a task from Pending to Running.
    pub fn start(&self, id: &str) {
        if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
            task.status = TaskStatus::Running;
        }
    }

    /// Update progress on a running task.
    pub fn update_progress(&self, id: &str, current: u64, total: u64, message: &str) {
        if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
            task.progress = Some(TaskProgress {
                current,
                total,
                message: message.to_string(),
            });
        }
    }

    /// Mark a task as completed with a result string.
    pub fn complete(&self, id: &str, result: String) {
        if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
            task.status = TaskStatus::Completed;
            task.progress = None;
            task.completed_at = Some(Utc::now().naive_utc());
            task.result = Some(result);
        }
    }

    /// Mark a task as failed with an error message.
    pub fn fail(&self, id: &str, error: String) {
        if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
            task.status = TaskStatus::Failed;
            task.progress = None;
            task.completed_at = Some(Utc::now().naive_utc());
            task.result = Some(error);
        }
    }

    /// Get a task by ID.
    pub fn get(&self, id: &str) -> Option<TaskInfo> {
        self.tasks.lock().unwrap().get(id).cloned()
    }

    /// List all tasks.
    pub fn list(&self) -> Vec<TaskInfo> {
        self.tasks.lock().unwrap().values().cloned().collect()
    }
}