use std::collections::BTreeMap;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum ComputeBackend {
    Cpu,
    Wasm,
    GpuExperimental,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct BackendCapability {
    pub backend: ComputeBackend,
    pub available: bool,
    pub supports_parallelism: bool,
    pub max_batch_size: Option<u32>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ComputeJob {
    pub job_id: Uuid,
    pub dataset_id: String,
    pub backend: ComputeBackend,
    pub operation: String,
    pub payload: Value,
    pub submitted_at: DateTime<Utc>,
    pub tags: BTreeMap<String, String>,
}

impl ComputeJob {
    pub fn new(
        dataset_id: impl Into<String>,
        backend: ComputeBackend,
        operation: impl Into<String>,
    ) -> Self {
        Self {
            job_id: Uuid::new_v4(),
            dataset_id: dataset_id.into(),
            backend,
            operation: operation.into(),
            payload: Value::Null,
            submitted_at: Utc::now(),
            tags: BTreeMap::new(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum ComputeStatus {
    Queued,
    Running,
    Completed,
    Failed,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct JobResult {
    pub job_id: Uuid,
    pub backend: ComputeBackend,
    pub status: ComputeStatus,
    pub output: Value,
    pub metrics: BTreeMap<String, f64>,
    pub finished_at: Option<DateTime<Utc>>,
}

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

    #[test]
    fn compute_job_new_sets_expected_defaults() {
        let job = ComputeJob::new("roads", ComputeBackend::Cpu, "buffer");

        assert_eq!(job.dataset_id, "roads");
        assert_eq!(job.backend, ComputeBackend::Cpu);
        assert_eq!(job.operation, "buffer");
        assert_eq!(job.payload, Value::Null);
        assert!(job.tags.is_empty());
    }

    #[test]
    fn job_result_round_trip_serialization() {
        let job = ComputeJob::new("roads", ComputeBackend::Wasm, "intersection");
        let result = JobResult {
            job_id: job.job_id,
            backend: job.backend,
            status: ComputeStatus::Completed,
            output: serde_json::json!({ "result": "ok" }),
            metrics: BTreeMap::from([(String::from("elapsed_ms"), 12.5)]),
            finished_at: Some(Utc::now()),
        };

        let json = serde_json::to_string(&result).expect("serialize result");
        let restored: JobResult = serde_json::from_str(&json).expect("deserialize result");

        assert_eq!(restored.status, ComputeStatus::Completed);
        assert_eq!(restored.metrics.get("elapsed_ms"), Some(&12.5));
    }
}
