pdk-unit 1.8.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use super::utils::{matches, not_found};
use crate::{Backend, UnitHttpMessage, UnitHttpRequest, UnitHttpResponse};
use regex::Regex;
use serde::Serialize;
use serde_json::json;
use std::cell::RefCell;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::HashMap;

const ETAG_HEADER: &str = "etag";
const IF_MATCH_HEADER: &str = "if-match";
const IF_NONE_HEADER: &str = "if-none-match";
const GET: &str = "GET";
const PUT: &str = "PUT";
const DELETE: &str = "DELETE";

thread_local! {
    static STORES: Regex = Regex::new(r"^/api/v1/stores$").unwrap();
    static CREATE_STORE: Regex = Regex::new(r"^/api/v1/stores/([^/]*)$").unwrap();
    static KEYS: Regex = Regex::new(r"^/api/v1/stores/([^/]*)/partitions/([^/]*)/keys$").unwrap();
    static KEY: Regex = Regex::new(r"^/api/v1/stores/([^/]*)/partitions/([^/]*)/keys/([^/]*)$").unwrap();
    static PARTITIONS: Regex = Regex::new(r"^/api/v1/stores/([^/]*)/partitions$").unwrap();
    static PARTITION: Regex = Regex::new(r"^/api/v1/stores/([^/]*)/partitions/([^/]*)$").unwrap();
}

type StoreStub = HashMap<String, HashMap<String, (Vec<u8>, String)>>;

#[derive(Debug, Default)]
pub struct OSBackend {
    stores: RefCell<HashMap<String, StoreStub>>,
}

impl Backend for OSBackend {
    fn call(&self, req: UnitHttpRequest) -> UnitHttpResponse {
        let path = req.header(":path").unwrap();
        let method = req.header(":method").unwrap();

        if let Some([store_id]) = matches(&CREATE_STORE, path) {
            if method == PUT {
                self.upsert_store(store_id)
            } else {
                not_found()
            }
        } else if let Some([]) = matches(&STORES, path) {
            if method == GET {
                self.get_stores()
            } else {
                not_found()
            }
        } else if let Some([store]) = matches(&PARTITIONS, path) {
            if method == GET {
                self.get_partitions(store)
            } else {
                not_found()
            }
        } else if let Some([store, partition]) = matches(&PARTITIONS, path) {
            if method == DELETE {
                self.delete_partition(store, partition)
            } else {
                not_found()
            }
        } else if let Some([store, partition]) = matches(&KEYS, path) {
            if method == GET {
                self.get_keys(store, partition)
            } else {
                not_found()
            }
        } else if let Some([store, partition, key]) = matches(&KEY, path) {
            if method == GET {
                self.get(store, partition, key)
            } else if method == PUT {
                self.store(
                    store,
                    partition,
                    key,
                    StoreMode::from(&req),
                    req.body().to_vec(),
                )
            } else if method == DELETE {
                self.delete(store, partition, key)
            } else {
                not_found()
            }
        } else {
            not_found()
        }
    }
}

impl OSBackend {
    fn upsert_store(&self, store: &str) -> UnitHttpResponse {
        self.stores
            .borrow_mut()
            .insert(store.to_string(), HashMap::new());

        ok(201)
    }

    fn get_stores(&self) -> UnitHttpResponse {
        let stores = Stores::new(
            self.stores
                .borrow()
                .keys()
                .map(|store_id| Store::new(store_id.to_string(), None))
                .collect::<Vec<Store>>(),
        );

        ok_with_body(stores)
    }

    fn get_keys(&self, store: &str, partition: &str) -> UnitHttpResponse {
        let keys = Keys::new(
            self.stores
                .borrow()
                .get(store)
                .and_then(|partitions| partitions.get(partition))
                .map(|hash| hash.keys().map(|key| Key::new(key.to_string())).collect())
                .unwrap_or_default(),
        );

        ok_with_body(keys)
    }

    fn get_partitions(&self, store: &str) -> UnitHttpResponse {
        let partitions = Partitions::new(
            self.stores
                .borrow()
                .get(store)
                .map(|hash| hash.keys().cloned().collect())
                .unwrap_or_default(),
        );

        ok_with_body(partitions)
    }

    fn store(
        &self,
        store: &str,
        partition: &str,
        key: &str,
        mode: StoreMode,
        item: Vec<u8>,
    ) -> UnitHttpResponse {
        let mut stores = self.stores.borrow_mut();
        let Some(store) = stores.get_mut(store) else {
            return store_not_found();
        };

        let partition = store
            .entry(partition.to_string())
            .or_insert_with(HashMap::new);

        let entry = partition.entry(key.to_string());

        let new_cas = match &entry {
            Occupied(occupied) => {
                let cas: u32 = occupied.get().1.parse().unwrap_or_default();
                (cas + 1).to_string()
            }
            Vacant(_) => "1".to_string(),
        };

        match mode {
            StoreMode::Always => {
                entry.insert_entry((item.to_vec(), new_cas));
            }
            StoreMode::Absent => {
                let Vacant(vacant) = entry else {
                    return cas_mismatch();
                };
                vacant.insert((item.to_vec(), new_cas));
            }
            StoreMode::Cas(cas) => {
                let Occupied(mut occupied) = entry else {
                    return cas_mismatch();
                };
                if occupied.get().1.eq(&cas) {
                    occupied.insert((item.to_vec(), new_cas));
                } else {
                    return cas_mismatch();
                }
            }
        }

        ok(201)
    }

    fn get(&self, store: &str, partition: &str, key: &str) -> UnitHttpResponse {
        let stores = self.stores.borrow();

        let Some(store) = stores.get(store) else {
            return store_not_found();
        };

        let Some(partition) = store.get(partition) else {
            return object_not_found();
        };

        let Some((obj, cas)) = partition.get(key) else {
            return object_not_found();
        };

        UnitHttpResponse::new(200)
            .with_header(":content-type", "application/json")
            .with_header(ETAG_HEADER, cas.as_str())
            .with_body(obj.clone())
    }

    fn delete(&self, store: &str, partition: &str, key: &str) -> UnitHttpResponse {
        if let Some(store) = self.stores.borrow_mut().get_mut(store) {
            if let Some(partition) = store.get_mut(partition) {
                partition.remove(key);
            };
        };

        ok(204)
    }

    fn delete_partition(&self, store: &str, partition: &str) -> UnitHttpResponse {
        if let Some(store) = self.stores.borrow_mut().get_mut(store) {
            store.remove(partition);
        };

        ok(204)
    }
}

#[derive(Serialize, Debug)]
pub struct Stores {
    pub values: Vec<Store>,
}

impl Stores {
    pub fn new(values: Vec<Store>) -> Self {
        Self { values }
    }
}

#[derive(Serialize, Debug)]
pub struct Store {
    #[serde(rename = "storeId")]
    store_id: String,

    #[serde(rename = "defaultTtlSeconds")]
    default_ttl_seconds: u32,
}

impl Store {
    pub fn new(store_id: String, ttl: Option<u32>) -> Self {
        Self {
            store_id,
            default_ttl_seconds: ttl.unwrap_or(u32::MAX),
        }
    }
}

#[derive(Serialize, Debug)]
pub struct Partitions {
    pub values: Vec<String>,
}

impl Partitions {
    pub fn new(values: Vec<String>) -> Self {
        Self { values }
    }
}

#[derive(Serialize, Debug)]
pub struct Keys {
    pub values: Vec<Key>,
}

impl Keys {
    pub fn new(values: Vec<Key>) -> Self {
        Self { values }
    }
}

#[derive(Serialize, Debug)]
pub struct Key {
    #[serde(rename = "keyId")]
    pub key_id: String,
}

impl Key {
    pub fn new(key_id: String) -> Self {
        Self { key_id }
    }
}

pub enum StoreMode {
    /// Indicates that the store operation does not depend on any condition.
    Always,
    /// Indicates that the store operation should succeed only if no other value was present in the store.
    Absent,
    /// Indicates that the store operation should succeed only if the stored value matches the provided cas.
    Cas(String),
}

impl StoreMode {
    pub fn from(request: &UnitHttpRequest) -> Self {
        let if_match = request.header(IF_MATCH_HEADER);
        let if_none = request.header(IF_NONE_HEADER);
        match (if_match, if_none) {
            (Some(value), None) => StoreMode::Cas(value.to_string()),
            (None, Some("*")) => StoreMode::Absent,
            (None, None) => StoreMode::Always,
            (_, _) => panic!("unexpected value"),
        }
    }
}

fn ok(status: u32) -> UnitHttpResponse {
    UnitHttpResponse::new(status)
}

fn ok_with_body<I: Serialize>(body: I) -> UnitHttpResponse {
    UnitHttpResponse::new(200)
        .with_header(":content-type", "application/json")
        .with_body(serde_json::to_vec(&body).unwrap_or_default())
}

fn store_not_found() -> UnitHttpResponse {
    UnitHttpResponse::new(404)
        .with_header(":content-type", "application/json")
        .with_body(json!({"message": "store not found"}).to_string())
}

fn object_not_found() -> UnitHttpResponse {
    UnitHttpResponse::new(404)
        .with_header(":content-type", "application/json")
        .with_body(json!({"message": "object with key not found"}).to_string())
}

fn cas_mismatch() -> UnitHttpResponse {
    UnitHttpResponse::new(412)
}