//! Cohort definition types: predicates, definitions, and the registry.
//!
//! A cohort is a named group of users defined by a predicate evaluated against
//! user metadata. The `CohortRegistry` stores all defined cohorts and provides
//! concurrent access via `DashMap`.

use std::sync::Arc;

use dashmap::{DashMap, mapref::entry::Entry};
use serde::{Deserialize, Serialize};

use crate::schema::TidalError;

/// Cheap-to-clone cohort name handle.
pub type CohortName = Arc<str>;

/// Maximum cohort-name length, in bytes (UTF-8 encoded).
///
/// The on-disk checkpoint suffix encodes the name length in a `u16`
/// (`cohort::checkpoint::cohort_entry_suffix`), so a name longer than
/// `u16::MAX` bytes cannot round-trip without truncating the length prefix and
/// corrupting the record. We validate against this bound at *define* time (here)
/// and again at *checkpoint* time (defense in depth for ad-hoc/restored names),
/// so an over-long name fails loudly with a typed error rather than silently
/// corrupting durable state.
pub const MAX_COHORT_NAME_LEN: usize = u16::MAX as usize;

/// A predicate evaluated against user metadata `HashMap<String, String>`.
///
/// Predicates compose via `And` and `Or` for arbitrary boolean expressions
/// over user attributes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Predicate {
    /// Exact match: `metadata[field] == value`.
    Eq { field: String, value: String },
    /// Set membership: `metadata[field]` is one of `values`.
    Any { field: String, values: Vec<String> },
    /// Numeric range (inclusive): `lo <= metadata[field].parse::<f64>() <= hi`.
    Range { field: String, lo: f64, hi: f64 },
    /// All sub-predicates must be true.
    And(Vec<Self>),
    /// At least one sub-predicate must be true.
    Or(Vec<Self>),
}

/// A named cohort definition consisting of a name and a predicate.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CohortDef {
    /// Human-readable cohort name (e.g., `"tech_enthusiasts"`).
    pub name: String,
    /// Predicate that determines cohort membership.
    pub predicate: Predicate,
}

/// Thread-safe registry of all defined cohorts.
///
/// Uses `DashMap` for concurrent reads/writes without mutex contention
/// on the hot path.
pub struct CohortRegistry {
    cohorts: DashMap<CohortName, CohortDef>,
}

impl CohortRegistry {
    /// Create an empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self {
            cohorts: DashMap::new(),
        }
    }

    /// Register a cohort definition.
    ///
    /// # Errors
    ///
    /// - `TidalError::InvalidInput` if the cohort name exceeds
    ///   [`MAX_COHORT_NAME_LEN`] bytes (it could not round-trip through a
    ///   checkpoint without truncation).
    /// - `TidalError::Internal` if a cohort with the same name already exists.
    pub fn define(&self, def: CohortDef) -> Result<(), TidalError> {
        // Validate length up front so an over-long name can never reach the
        // checkpoint suffix encoder, where it would silently truncate.
        if def.name.len() > MAX_COHORT_NAME_LEN {
            return Err(TidalError::invalid_input(format!(
                "cohort name is {} bytes, exceeds maximum of {MAX_COHORT_NAME_LEN} bytes",
                def.name.len(),
            )));
        }

        let name: CohortName = Arc::from(def.name.as_str());
        // Atomic check-and-insert via the entry API. A bare
        // `contains_key`/`insert` pair has a TOCTOU window where two concurrent
        // `define` calls for the same name both see "absent" and the second
        // silently overwrites the first (dropping the duplicate error). Holding
        // the entry's bucket lock across the check and the insert closes that
        // window: exactly one caller observes `Vacant` and the other gets the
        // duplicate error.
        match self.cohorts.entry(name) {
            Entry::Occupied(_) => Err(TidalError::internal(
                "define_cohort",
                format!("cohort '{}' already defined", def.name),
            )),
            Entry::Vacant(slot) => {
                slot.insert(def);
                Ok(())
            }
        }
    }

    /// Look up a cohort by name.
    ///
    /// Returns a `DashMap` reference guard or `None` if not found.
    ///
    /// The lookup probes by borrowed `&str`: `Arc<str>: Borrow<str>`, so
    /// `DashMap::get` hashes and compares the key in place with no per-probe
    /// allocation. (The previous `Arc::from(name)` allocated a throwaway
    /// `Arc<str>` on every call — including the resolver's per-cohort probes in
    /// its membership-resolution loop — purely to satisfy the key type.)
    #[must_use]
    pub fn get(&self, name: &str) -> Option<dashmap::mapref::one::Ref<'_, CohortName, CohortDef>> {
        self.cohorts.get(name)
    }

    /// List all registered cohort names.
    #[must_use]
    pub fn names(&self) -> Vec<CohortName> {
        self.cohorts.iter().map(|r| r.key().clone()).collect()
    }

    /// Number of registered cohorts.
    #[must_use]
    pub fn len(&self) -> usize {
        self.cohorts.len()
    }

    /// Whether the registry is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.cohorts.is_empty()
    }
}

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

impl std::fmt::Debug for CohortRegistry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CohortRegistry")
            .field("count", &self.cohorts.len())
            .finish_non_exhaustive()
    }
}

// ── Serialization ─────────────────────────────────────────────────────────────

/// Serialize a cohort definition as JSON bytes for durable storage.
///
/// # Errors
///
/// Returns `TidalError::Internal` if serialization fails. `serde_json::to_vec`
/// cannot fail for the current `CohortDef`/`Predicate` shape (no maps with
/// non-string keys, no recursion-limit overflow), so this is effectively
/// infallible today. We propagate the error rather than swallowing it because
/// the previous `unwrap_or_default()` turned the impossible failure into an
/// empty `Vec` — and an empty `Vec` persisted to `Tag::CohortDef` deserializes
/// back to `None` on the next open, silently dropping the cohort definition
/// after a restart. Surfacing the error keeps that footgun out of the durable
/// write path if a future `Predicate` variant ever makes serialization fallible.
pub fn serialize_cohort_def(def: &CohortDef) -> Result<Vec<u8>, TidalError> {
    serde_json::to_vec(def).map_err(|e| {
        TidalError::internal(
            "serialize_cohort_def",
            format!("failed to serialize cohort '{}': {e}", def.name),
        )
    })
}

/// Deserialize a cohort definition from JSON bytes.
///
/// Returns `None` on malformed input -- never panics.
#[must_use]
pub fn deserialize_cohort_def(bytes: &[u8]) -> Option<CohortDef> {
    serde_json::from_slice(bytes).ok()
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    fn tech_predicate() -> Predicate {
        Predicate::And(vec![
            Predicate::Eq {
                field: "locale".into(),
                value: "en".into(),
            },
            Predicate::Eq {
                field: "primary_category".into(),
                value: "tech".into(),
            },
        ])
    }

    #[test]
    fn define_and_get() {
        let registry = CohortRegistry::new();
        let def = CohortDef {
            name: "tech".into(),
            predicate: tech_predicate(),
        };
        registry.define(def).unwrap();
        assert!(registry.get("tech").is_some());
        assert!(registry.get("nonexistent").is_none());
    }

    #[test]
    fn define_duplicate_fails() {
        let registry = CohortRegistry::new();
        let def1 = CohortDef {
            name: "tech".into(),
            predicate: tech_predicate(),
        };
        let def2 = CohortDef {
            name: "tech".into(),
            predicate: tech_predicate(),
        };
        registry.define(def1).unwrap();
        assert!(registry.define(def2).is_err());
    }

    #[test]
    fn define_rejects_over_long_name() {
        let registry = CohortRegistry::new();
        let def = CohortDef {
            name: "n".repeat(MAX_COHORT_NAME_LEN + 1),
            predicate: tech_predicate(),
        };
        let err = registry
            .define(def)
            .expect_err("over-long name must be rejected");
        assert!(matches!(err, TidalError::InvalidInput(_)));
        assert!(registry.is_empty());
    }

    #[test]
    fn concurrent_define_same_name_yields_exactly_one_winner() {
        use std::sync::{
            Arc as StdArc,
            atomic::{AtomicUsize, Ordering},
        };
        use std::thread;

        // Hammer the same cohort name from many threads. The atomic entry-API
        // check-and-insert must admit exactly one definition; every other call
        // must observe the duplicate error. A TOCTOU race would let >1 succeed
        // (the loser silently overwriting the winner) — this asserts it cannot.
        let registry = StdArc::new(CohortRegistry::new());
        let barrier = StdArc::new(std::sync::Barrier::new(16));
        let successes = StdArc::new(AtomicUsize::new(0));

        let handles: Vec<_> = (0..16)
            .map(|_| {
                let registry = StdArc::clone(&registry);
                let barrier = StdArc::clone(&barrier);
                let successes = StdArc::clone(&successes);
                thread::spawn(move || {
                    let def = CohortDef {
                        name: "contended".into(),
                        predicate: tech_predicate(),
                    };
                    barrier.wait();
                    if registry.define(def).is_ok() {
                        successes.fetch_add(1, Ordering::Relaxed);
                    }
                })
            })
            .collect();

        for h in handles {
            h.join().expect("thread did not panic");
        }

        assert_eq!(
            successes.load(Ordering::Relaxed),
            1,
            "exactly one concurrent define must win"
        );
        assert_eq!(registry.len(), 1);
    }

    #[test]
    fn names_and_len() {
        let registry = CohortRegistry::new();
        assert!(registry.is_empty());
        assert_eq!(registry.len(), 0);

        registry
            .define(CohortDef {
                name: "a".into(),
                predicate: Predicate::Eq {
                    field: "x".into(),
                    value: "1".into(),
                },
            })
            .unwrap();
        registry
            .define(CohortDef {
                name: "b".into(),
                predicate: Predicate::Eq {
                    field: "x".into(),
                    value: "2".into(),
                },
            })
            .unwrap();

        assert_eq!(registry.len(), 2);
        assert!(!registry.is_empty());

        let mut names: Vec<String> = registry.names().iter().map(|n| n.to_string()).collect();
        names.sort();
        assert_eq!(names, vec!["a", "b"]);
    }

    #[test]
    fn registry_is_send_and_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<CohortRegistry>();
    }

    #[test]
    fn serde_roundtrip_simple() {
        let def = CohortDef {
            name: "tech_en".into(),
            predicate: Predicate::Eq {
                field: "locale".into(),
                value: "en".into(),
            },
        };
        let bytes = serialize_cohort_def(&def).unwrap();
        let restored = deserialize_cohort_def(&bytes).unwrap();
        assert_eq!(restored.name, "tech_en");
        match &restored.predicate {
            Predicate::Eq { field, value } => {
                assert_eq!(field, "locale");
                assert_eq!(value, "en");
            }
            other => panic!("expected Eq, got {other:?}"),
        }
    }

    #[test]
    fn serde_roundtrip_compound() {
        let def = CohortDef {
            name: "compound".into(),
            predicate: Predicate::And(vec![
                Predicate::Eq {
                    field: "locale".into(),
                    value: "en".into(),
                },
                Predicate::Or(vec![
                    Predicate::Any {
                        field: "interest".into(),
                        values: vec!["tech".into(), "science".into()],
                    },
                    Predicate::Range {
                        field: "age".into(),
                        lo: 18.0,
                        hi: 65.0,
                    },
                ]),
            ]),
        };
        let bytes = serialize_cohort_def(&def).unwrap();
        let restored = deserialize_cohort_def(&bytes).unwrap();
        assert_eq!(restored.name, "compound");
        // Verify the structure survived.
        match &restored.predicate {
            Predicate::And(parts) => assert_eq!(parts.len(), 2),
            other => panic!("expected And, got {other:?}"),
        }
    }

    #[test]
    fn deserialize_malformed_returns_none() {
        assert!(deserialize_cohort_def(b"").is_none());
        assert!(deserialize_cohort_def(b"not json").is_none());
        assert!(deserialize_cohort_def(b"{}").is_none());
    }
}
