//! Transaction Mode Selector
//!
//! Handles the user's explicit choice between public and private
//! transaction modes. The mode is determined BEFORE signing and
//! cannot be changed after proof generation.
//!
//! # Security
//!
//! - Mode must be explicitly selected, never defaulted
//! - Mode is stored in the transaction context
//! - Mode switching after proof creation is forbidden
//! - The signer must reject ambiguous mode states
//!
//! # State Machine
//!
//! ```text
//! UNSELECTED ──► SELECTED ──► LOCKED
//!     │              │
//!     └──────────────┴──► REJECTED (terminal)
//! ```

use coldstar_zk::types::TransactionMode;
use serde::{Deserialize, Serialize};
use std::fmt;
use thiserror::Error;

/// Errors from mode selection operations.
#[derive(Error, Debug)]
pub enum ModeError {
    /// Attempted to change mode after it was locked.
    #[error("cannot change mode after proof generation — start a new transaction to change modes")]
    ModeLocked,

    /// Attempted operation on a rejected mode.
    #[error("transaction mode was rejected — start a new transaction")]
    ModeRejected,

    /// Invalid mode string provided.
    #[error("invalid mode: '{0}' — must be 'public' or 'private'")]
    InvalidMode(String),

    /// Attempted to lock without a selected mode.
    #[error("cannot lock — no mode has been selected")]
    NoModeSelected,

    /// Mode does not match expected value.
    #[error("mode mismatch: expected {expected}, got {actual}")]
    ModeMismatch {
        expected: String,
        actual: String,
    },
}

/// State machine for transaction mode lifecycle.
///
/// Transitions:
///   - `Unselected` -> `Selected` -> `Locked`
///   - (Any state) -> `Rejected` (terminal)
///
/// Once `Locked`, the mode cannot be changed without starting a new transaction.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ModeState {
    /// No mode chosen yet.
    Unselected,
    /// User has chosen a mode.
    Selected,
    /// Proofs generated / signing prep done — no changes allowed.
    Locked,
    /// Mode was rejected (e.g., invalid transition). Terminal state.
    Rejected,
}

impl fmt::Display for ModeState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ModeState::Unselected => write!(f, "unselected"),
            ModeState::Selected => write!(f, "selected"),
            ModeState::Locked => write!(f, "locked"),
            ModeState::Rejected => write!(f, "rejected"),
        }
    }
}

/// The result of a mode selection operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModeSelection {
    /// The selected transaction mode, if any.
    pub mode: Option<TransactionMode>,
    /// Current state of the mode selector.
    pub state: ModeState,
    /// Human-readable reason for the current state.
    #[serde(default)]
    pub reason: String,
}

impl ModeSelection {
    /// Returns `true` if a valid mode has been selected or locked.
    pub fn is_valid(&self) -> bool {
        self.mode.is_some()
            && (self.state == ModeState::Selected || self.state == ModeState::Locked)
    }

    /// Returns `true` if the mode is locked (no further changes possible).
    pub fn is_locked(&self) -> bool {
        self.state == ModeState::Locked
    }
}

/// Transaction mode selector.
///
/// Manages the lifecycle of mode selection for a single transaction.
///
/// # Example
///
/// ```
/// use coldstar_privacy::ModeSelector;
///
/// let mut selector = ModeSelector::new();
/// selector.select("private").unwrap();
/// selector.lock().unwrap();
/// // selector.select("public") would now return Err(ModeError::ModeLocked)
/// ```
pub struct ModeSelector {
    mode: Option<TransactionMode>,
    state: ModeState,
    reason: String,
}

impl ModeSelector {
    /// Create a new mode selector in the `Unselected` state.
    pub fn new() -> Self {
        Self {
            mode: None,
            state: ModeState::Unselected,
            reason: String::new(),
        }
    }

    /// Returns the currently selected mode, if any.
    pub fn mode(&self) -> Option<TransactionMode> {
        self.mode
    }

    /// Returns the current state of the selector.
    pub fn state(&self) -> ModeState {
        self.state
    }

    /// Returns a snapshot of the current selection state.
    pub fn selection(&self) -> ModeSelection {
        ModeSelection {
            mode: self.mode,
            state: self.state,
            reason: self.reason.clone(),
        }
    }

    /// Select a transaction mode.
    ///
    /// # Arguments
    ///
    /// * `mode_str` - `"public"` or `"private"`
    ///
    /// # Errors
    ///
    /// Returns [`ModeError`] if the state transition is invalid or the mode
    /// string is unrecognized.
    pub fn select(&mut self, mode_str: &str) -> Result<ModeSelection, ModeError> {
        if self.state == ModeState::Locked {
            return Err(ModeError::ModeLocked);
        }

        if self.state == ModeState::Rejected {
            return Err(ModeError::ModeRejected);
        }

        match TransactionMode::from_str_strict(mode_str) {
            Some(mode) => {
                self.mode = Some(mode);
                self.state = ModeState::Selected;
                self.reason = format!("Mode set to {}", mode.as_str());
                Ok(self.selection())
            }
            None => {
                self.state = ModeState::Rejected;
                self.reason = format!(
                    "Invalid mode: '{}'. Must be 'public' or 'private'.",
                    mode_str
                );
                Err(ModeError::InvalidMode(mode_str.to_string()))
            }
        }
    }

    /// Lock the mode — prevents further changes.
    ///
    /// Should be called after proof generation (private) or
    /// before signing (public).
    ///
    /// # Errors
    ///
    /// Returns [`ModeError`] if no mode has been selected or mode was rejected.
    pub fn lock(&mut self) -> Result<ModeSelection, ModeError> {
        match self.state {
            ModeState::Unselected => Err(ModeError::NoModeSelected),
            ModeState::Rejected => Err(ModeError::ModeRejected),
            ModeState::Selected | ModeState::Locked => {
                self.state = ModeState::Locked;
                if let Some(mode) = self.mode {
                    self.reason = format!("Mode locked to {}", mode.as_str());
                }
                Ok(self.selection())
            }
        }
    }

    /// Reset the selector for a new transaction.
    pub fn reset(&mut self) {
        self.mode = None;
        self.state = ModeState::Unselected;
        self.reason = String::new();
    }

    /// Check that the current mode matches the expected mode.
    ///
    /// # Errors
    ///
    /// Returns [`ModeError`] if no mode is selected or the mode doesn't match.
    pub fn require_mode(&self, expected: TransactionMode) -> Result<(), ModeError> {
        match self.mode {
            None => Err(ModeError::NoModeSelected),
            Some(actual) if actual != expected => Err(ModeError::ModeMismatch {
                expected: expected.as_str().to_string(),
                actual: actual.as_str().to_string(),
            }),
            _ => Ok(()),
        }
    }

    /// Get a human-readable status string.
    pub fn display_status(&self) -> String {
        match self.state {
            ModeState::Unselected => "Mode: NOT SELECTED".to_string(),
            ModeState::Selected => {
                if let Some(mode) = self.mode {
                    format!("Mode: {} (selected, not locked)", mode.as_str().to_uppercase())
                } else {
                    "Mode: SELECTED (unknown)".to_string()
                }
            }
            ModeState::Locked => {
                if let Some(mode) = self.mode {
                    format!("Mode: {} (LOCKED)", mode.as_str().to_uppercase())
                } else {
                    "Mode: LOCKED (unknown)".to_string()
                }
            }
            ModeState::Rejected => {
                format!("Mode: REJECTED ({})", self.reason)
            }
        }
    }
}

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

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

    #[test]
    fn new_selector_is_unselected() {
        let selector = ModeSelector::new();
        assert_eq!(selector.state(), ModeState::Unselected);
        assert!(selector.mode().is_none());
    }

    #[test]
    fn select_public() {
        let mut selector = ModeSelector::new();
        let result = selector.select("public").unwrap();
        assert_eq!(result.mode, Some(TransactionMode::Public));
        assert_eq!(result.state, ModeState::Selected);
        assert!(result.is_valid());
        assert!(!result.is_locked());
    }

    #[test]
    fn select_private() {
        let mut selector = ModeSelector::new();
        let result = selector.select("private").unwrap();
        assert_eq!(result.mode, Some(TransactionMode::Private));
        assert_eq!(result.state, ModeState::Selected);
    }

    #[test]
    fn lock_after_select() {
        let mut selector = ModeSelector::new();
        selector.select("private").unwrap();
        let result = selector.lock().unwrap();
        assert!(result.is_locked());
        assert!(result.is_valid());
    }

    #[test]
    fn cannot_select_after_lock() {
        let mut selector = ModeSelector::new();
        selector.select("private").unwrap();
        selector.lock().unwrap();
        assert!(matches!(
            selector.select("public"),
            Err(ModeError::ModeLocked)
        ));
    }

    #[test]
    fn cannot_lock_without_selection() {
        let mut selector = ModeSelector::new();
        assert!(matches!(selector.lock(), Err(ModeError::NoModeSelected)));
    }

    #[test]
    fn invalid_mode_rejects() {
        let mut selector = ModeSelector::new();
        assert!(selector.select("invalid").is_err());
        assert_eq!(selector.state(), ModeState::Rejected);
    }

    #[test]
    fn cannot_select_after_reject() {
        let mut selector = ModeSelector::new();
        let _ = selector.select("bad");
        assert!(matches!(
            selector.select("public"),
            Err(ModeError::ModeRejected)
        ));
    }

    #[test]
    fn reset_clears_state() {
        let mut selector = ModeSelector::new();
        selector.select("private").unwrap();
        selector.lock().unwrap();
        selector.reset();
        assert_eq!(selector.state(), ModeState::Unselected);
        assert!(selector.mode().is_none());
        // Can select again after reset
        selector.select("public").unwrap();
        assert_eq!(selector.mode(), Some(TransactionMode::Public));
    }

    #[test]
    fn require_mode_matches() {
        let mut selector = ModeSelector::new();
        selector.select("private").unwrap();
        assert!(selector.require_mode(TransactionMode::Private).is_ok());
    }

    #[test]
    fn require_mode_mismatch() {
        let mut selector = ModeSelector::new();
        selector.select("private").unwrap();
        assert!(matches!(
            selector.require_mode(TransactionMode::Public),
            Err(ModeError::ModeMismatch { .. })
        ));
    }

    #[test]
    fn display_status_shows_state() {
        let mut selector = ModeSelector::new();
        assert!(selector.display_status().contains("NOT SELECTED"));

        selector.select("private").unwrap();
        assert!(selector.display_status().contains("PRIVATE"));
        assert!(selector.display_status().contains("not locked"));

        selector.lock().unwrap();
        assert!(selector.display_status().contains("LOCKED"));
    }
}
