//! Online property monitoring for simulation runs.

use std::collections::{BTreeMap, VecDeque};
use telltale_types::FixedQ32;

use telltale_machine::coroutine::{CoroStatus, Coroutine};
use telltale_machine::instr::Endpoint;
use telltale_machine::model::state::{SessionId, SessionState};
use telltale_machine::trace::obs_session;
use telltale_machine::ObsEvent;
use telltale_types::LocalTypeR;

use crate::value_conv::registers_to_f64s;

/// Tolerance for fixed-point comparisons (approximately 1e-8).
fn tolerance() -> FixedQ32 {
    FixedQ32::from_ratio(1, 100_000_000).expect("tolerance must be representable")
}

/// Property check context (read-only).
pub struct PropertyContext<'a> {
    /// Current simulation tick.
    pub tick: u64,
    /// Observable event trace from the ProtocolMachine.
    pub trace: &'a [ObsEvent],
    /// Session snapshots for buffer and type state inspection.
    pub sessions: &'a BTreeMap<SessionId, SessionState>,
    /// Coroutine snapshots for register inspection.
    pub coroutines: &'a [Coroutine],
}

/// Built-in properties supported by the monitor.
#[derive(Debug, Clone)]
pub enum Property {
    /// No coroutines have faulted.
    NoFaults,
    /// All coroutine states remain on the probability simplex.
    Simplex,
    /// Every send is received within a bounded number of ticks.
    SendRecvLiveness {
        /// Session to monitor.
        sid: SessionId,
        /// Maximum ticks between send and receive.
        bound: usize,
    },
    /// Session type depth monotonically decreases.
    TypeMonotonicity {
        /// Session to monitor.
        sid: SessionId,
    },
    /// Buffer occupancy stays below a maximum.
    BufferBound {
        /// Session to monitor.
        sid: SessionId,
        /// Maximum buffer size.
        max: usize,
    },
    /// Custom liveness property with precondition and goal predicates.
    Liveness {
        /// Property name for reporting.
        name: String,
        /// Predicate that starts the liveness window.
        precondition: Predicate,
        /// Predicate that must become true within bound.
        goal: Predicate,
        /// Maximum steps to reach goal after precondition.
        bound: usize,
    },
}

impl Property {
    /// Returns the property name for reporting.
    #[must_use]
    pub fn name(&self) -> String {
        match self {
            Self::NoFaults => "no_faults".into(),
            Self::Simplex => "simplex".into(),
            Self::SendRecvLiveness { sid, .. } => format!("send_recv_liveness[{sid}]"),
            Self::TypeMonotonicity { sid } => format!("type_monotonicity[{sid}]"),
            Self::BufferBound { sid, max } => format!("buffer_bound[{sid},{max}]"),
            Self::Liveness { name, .. } => name.clone(),
        }
    }
}

/// Predicate for liveness properties.
#[derive(Debug, Clone)]
pub enum Predicate {
    /// A built-in predicate like no_faults or simplex.
    BuiltIn(BuiltInPredicate),
    /// Compare current tick against a threshold.
    TickCmp {
        /// Comparison operator.
        op: CmpOp,
        /// Right-hand side value.
        rhs: u64,
    },
    /// Compare distance to equilibrium against a threshold.
    DistanceToEquilibriumCmp {
        /// Comparison operator.
        op: CmpOp,
        /// Right-hand side value.
        rhs: FixedQ32,
    },
}

/// Built-in predicates for property evaluation.
#[derive(Debug, Clone)]
pub enum BuiltInPredicate {
    /// True when no coroutines have faulted.
    NoFaults,
    /// True when all states remain on the probability simplex.
    Simplex,
}

/// Comparison operators for predicate evaluation.
#[derive(Debug, Clone, Copy)]
pub enum CmpOp {
    /// Less than.
    Lt,
    /// Less than or equal.
    Lte,
    /// Greater than.
    Gt,
    /// Greater than or equal.
    Gte,
    /// Equal.
    Eq,
}

impl Predicate {
    /// Evaluates the predicate against the current context.
    #[must_use]
    pub fn eval(&self, ctx: &PropertyContext) -> bool {
        match self {
            Self::BuiltIn(b) => match b {
                BuiltInPredicate::NoFaults => check_no_faults(ctx),
                BuiltInPredicate::Simplex => check_simplex(ctx),
            },
            Self::TickCmp { op, rhs } => match op {
                CmpOp::Lt => ctx.tick < *rhs,
                CmpOp::Lte => ctx.tick <= *rhs,
                CmpOp::Gt => ctx.tick > *rhs,
                CmpOp::Gte => ctx.tick >= *rhs,
                CmpOp::Eq => ctx.tick == *rhs,
            },
            Self::DistanceToEquilibriumCmp { op, rhs } => {
                let Some(dist) = distance_to_equilibrium(ctx) else {
                    return false;
                };
                match op {
                    CmpOp::Lt => dist < *rhs,
                    CmpOp::Lte => dist <= *rhs,
                    CmpOp::Gt => dist > *rhs,
                    CmpOp::Gte => dist >= *rhs,
                    CmpOp::Eq => (dist - *rhs).abs() < tolerance(),
                }
            }
        }
    }
}

/// Parse an invariant/property declaration into the canonical property model.
pub fn parse_property(s: &str) -> Result<Property, String> {
    let name = s.trim();
    match name {
        "no_faults" => return Ok(Property::NoFaults),
        "simplex" => return Ok(Property::Simplex),
        _ => {}
    }

    if let Some((base, args)) = parse_call(name) {
        match base {
            "send_recv_liveness" => return parse_send_recv_liveness(&args),
            "buffer_bound" => return parse_buffer_bound(&args),
            "type_monotonicity" => return parse_type_monotonicity(&args),
            _ => {}
        }
    }

    Err(format!("unknown property: {name}"))
}

/// Parse a simple predicate expression.
///
/// Supported forms:
/// - `no_faults`
/// - `simplex`
/// - `tick < N`, `tick <= N`, `tick > N`, `tick >= N`, `tick == N`
/// - `distance_to_equilibrium < X` (compared against uniform equilibrium)
pub fn parse_predicate(s: &str) -> Result<Predicate, String> {
    let expr = s.trim();
    match expr {
        "no_faults" => return Ok(Predicate::BuiltIn(BuiltInPredicate::NoFaults)),
        "simplex" => return Ok(Predicate::BuiltIn(BuiltInPredicate::Simplex)),
        _ => {}
    }

    let ops = [">=", "<=", "==", ">", "<"];
    for op in ops {
        if let Some((lhs, rhs)) = expr.split_once(op) {
            let lhs = lhs.trim();
            let rhs = rhs.trim();
            let op = match op {
                ">" => CmpOp::Gt,
                ">=" => CmpOp::Gte,
                "<" => CmpOp::Lt,
                "<=" => CmpOp::Lte,
                "==" => CmpOp::Eq,
                _ => return Err(format!("unsupported operator: {op}")),
            };
            if lhs == "tick" {
                let rhs = rhs
                    .parse::<u64>()
                    .map_err(|_| format!("invalid numeric literal: {rhs}"))?;
                return Ok(Predicate::TickCmp { op, rhs });
            }
            if lhs == "distance_to_equilibrium" {
                let rhs = rhs
                    .parse::<FixedQ32>()
                    .map_err(|_| format!("invalid numeric literal: {rhs}"))?;
                return Ok(Predicate::DistanceToEquilibriumCmp { op, rhs });
            }
            return Err(format!("unsupported predicate lhs: {lhs}"));
        }
    }

    Err(format!("unsupported predicate: {expr}"))
}

fn parse_send_recv_liveness(args: &[&str]) -> Result<Property, String> {
    if args.len() != 2 {
        return Err("send_recv_liveness requires (sid,bound)".into());
    }
    let sid = args[0]
        .parse::<usize>()
        .map_err(|_| "invalid session id".to_string())?;
    let bound = args[1]
        .parse::<usize>()
        .map_err(|_| "invalid bound".to_string())?;
    Ok(Property::SendRecvLiveness { sid, bound })
}

fn parse_buffer_bound(args: &[&str]) -> Result<Property, String> {
    if args.len() != 2 {
        return Err("buffer_bound requires (sid,max)".into());
    }
    let sid = args[0]
        .parse::<usize>()
        .map_err(|_| "invalid session id".to_string())?;
    let max = args[1]
        .parse::<usize>()
        .map_err(|_| "invalid max".to_string())?;
    Ok(Property::BufferBound { sid, max })
}

fn parse_type_monotonicity(args: &[&str]) -> Result<Property, String> {
    if args.len() != 1 {
        return Err("type_monotonicity requires (sid)".into());
    }
    let sid = args[0]
        .parse::<usize>()
        .map_err(|_| "invalid session id".to_string())?;
    Ok(Property::TypeMonotonicity { sid })
}

fn parse_call(expr: &str) -> Option<(&str, Vec<&str>)> {
    let expr = expr.trim();
    let expr = expr.strip_suffix(')')?;
    let (name, args) = expr.split_once('(')?;
    let args = args
        .split(',')
        .map(|s| s.trim())
        .filter(|s| !s.is_empty())
        .collect::<Vec<_>>();
    Some((name.trim(), args))
}

/// Property violation record.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct PropertyViolation {
    /// Name of the violated property.
    pub property_name: String,
    /// Simulation tick when the violation occurred.
    pub tick: u64,
    /// Human-readable violation description.
    pub details: String,
}

struct LivenessTracking {
    precondition_tick: Option<u64>,
    steps_since_precondition: usize,
}

struct SendRecvTracking {
    pending_sends: BTreeMap<(String, String, String), VecDeque<u64>>,
}

struct TypeMonotonicityTracking {
    last_depths: BTreeMap<Endpoint, usize>,
}

enum PropertyState {
    Invariant {
        violated: bool,
    },
    Liveness {
        violated: bool,
        tracking: LivenessTracking,
    },
    SendRecv {
        violated: bool,
        tracking: SendRecvTracking,
    },
    TypeMonotonicity {
        violated: bool,
        tracking: TypeMonotonicityTracking,
    },
}

/// Online property monitor. Hooks into the ProtocolMachine step loop.
pub struct PropertyMonitor {
    properties: Vec<Property>,
    states: Vec<PropertyState>,
    violations: Vec<PropertyViolation>,
    last_trace_len: usize,
    session_counters: BTreeMap<SessionId, u64>,
}

impl PropertyMonitor {
    /// Create a new monitor for a set of properties.
    #[must_use]
    pub fn new(properties: Vec<Property>) -> Self {
        let states = properties
            .iter()
            .map(|p| match p {
                Property::SendRecvLiveness { .. } => PropertyState::SendRecv {
                    violated: false,
                    tracking: SendRecvTracking {
                        pending_sends: BTreeMap::new(),
                    },
                },
                Property::TypeMonotonicity { .. } => PropertyState::TypeMonotonicity {
                    violated: false,
                    tracking: TypeMonotonicityTracking {
                        last_depths: BTreeMap::new(),
                    },
                },
                Property::Liveness { .. } => PropertyState::Liveness {
                    violated: false,
                    tracking: LivenessTracking {
                        precondition_tick: None,
                        steps_since_precondition: 0,
                    },
                },
                _ => PropertyState::Invariant { violated: false },
            })
            .collect();

        Self {
            properties,
            states,
            violations: Vec::new(),
            last_trace_len: 0,
            session_counters: BTreeMap::new(),
        }
    }

    /// Access recorded violations.
    #[must_use]
    pub fn violations(&self) -> &[PropertyViolation] {
        &self.violations
    }

    /// Check all properties against the current ProtocolMachine context.
    #[allow(clippy::too_many_lines)]
    pub fn check(&mut self, ctx: &PropertyContext<'_>) {
        let new_events = ctx.trace.get(self.last_trace_len..).unwrap_or(&[]);
        let mut new_events_local = Vec::with_capacity(new_events.len());
        for event in new_events {
            let local_tick = obs_session(event).map(|sid| {
                let counter = self.session_counters.entry(sid).or_insert(0);
                let tick = *counter;
                *counter += 1;
                tick
            });
            new_events_local.push((event, local_tick));
        }

        for (idx, prop) in self.properties.iter().enumerate() {
            match (&prop, &mut self.states[idx]) {
                (Property::NoFaults, PropertyState::Invariant { violated }) => {
                    if !*violated && !check_no_faults(ctx) {
                        *violated = true;
                        self.violations.push(PropertyViolation {
                            property_name: prop.name(),
                            tick: ctx.tick,
                            details: "faulted coroutine detected".into(),
                        });
                    }
                }
                (Property::Simplex, PropertyState::Invariant { violated }) => {
                    if !*violated && !check_simplex(ctx) {
                        *violated = true;
                        self.violations.push(PropertyViolation {
                            property_name: prop.name(),
                            tick: ctx.tick,
                            details: "state left simplex".into(),
                        });
                    }
                }
                (Property::BufferBound { sid, max }, PropertyState::Invariant { violated }) => {
                    if *violated {
                        continue;
                    }
                    if let Some(session) = ctx.sessions.get(sid) {
                        let too_large = session.buffers.values().any(|buf| buf.len() > *max);
                        if too_large {
                            *violated = true;
                            self.violations.push(PropertyViolation {
                                property_name: prop.name(),
                                tick: ctx.tick,
                                details: format!("buffer size exceeded {max}"),
                            });
                        }
                    }
                }
                (
                    Property::SendRecvLiveness { sid, bound },
                    PropertyState::SendRecv { violated, tracking },
                ) => {
                    if *violated {
                        continue;
                    }
                    for (event, local_tick) in &new_events_local {
                        match event {
                            ObsEvent::Sent {
                                session,
                                from,
                                to,
                                label,
                                ..
                            } if session == sid => {
                                if let Some(tick) = *local_tick {
                                    let key = (from.clone(), to.clone(), label.clone());
                                    tracking
                                        .pending_sends
                                        .entry(key)
                                        .or_default()
                                        .push_back(tick);
                                }
                            }
                            ObsEvent::Received {
                                session,
                                from,
                                to,
                                label,
                                ..
                            } if session == sid => {
                                let key = (from.clone(), to.clone(), label.clone());
                                if let Some(queue) = tracking.pending_sends.get_mut(&key) {
                                    queue.pop_front();
                                    if queue.is_empty() {
                                        tracking.pending_sends.remove(&key);
                                    }
                                }
                            }
                            _ => {}
                        }
                    }
                    let oldest_send = tracking.pending_sends.iter().fold(
                        None,
                        |oldest, (key, queue)| match queue.front() {
                            Some(local_tick) => match oldest {
                                Some((_, oldest_tick)) if oldest_tick <= *local_tick => oldest,
                                _ => Some((key, *local_tick)),
                            },
                            None => oldest,
                        },
                    );
                    if let Some(((from, to, label), local_tick)) = oldest_send {
                        let current_tick = self.session_counters.get(sid).copied().unwrap_or(0);
                        let bound_u64 = u64::try_from(*bound).unwrap_or(u64::MAX);
                        if current_tick.saturating_sub(local_tick) > bound_u64 {
                            *violated = true;
                            self.violations.push(PropertyViolation {
                                property_name: prop.name(),
                                tick: ctx.tick,
                                details: format!(
                                    "send {} -> {} [{}] at local tick {} not received within {bound} ticks",
                                    from, to, label, local_tick
                                ),
                            });
                        }
                    }
                }
                (
                    Property::TypeMonotonicity { sid },
                    PropertyState::TypeMonotonicity { violated, tracking },
                ) => {
                    if *violated {
                        continue;
                    }
                    if let Some(session) = ctx.sessions.get(sid) {
                        for (ep, entry) in &session.local_types {
                            if contains_recursion(&entry.original) {
                                continue;
                            }
                            let depth = type_depth(&entry.current);
                            if let Some(prev) = tracking.last_depths.get(ep) {
                                if depth > *prev {
                                    *violated = true;
                                    self.violations.push(PropertyViolation {
                                        property_name: prop.name(),
                                        tick: ctx.tick,
                                        details: format!(
                                            "type depth increased for {}:{} ({prev} -> {depth})",
                                            ep.sid, ep.role
                                        ),
                                    });
                                    break;
                                }
                            }
                            tracking.last_depths.insert(ep.clone(), depth);
                        }
                    }
                }
                (
                    Property::Liveness {
                        precondition,
                        goal,
                        bound,
                        ..
                    },
                    PropertyState::Liveness { violated, tracking },
                ) => {
                    if *violated {
                        continue;
                    }
                    let precond = precondition.eval(ctx);
                    if tracking.precondition_tick.is_none() && precond {
                        tracking.precondition_tick = Some(ctx.tick);
                        tracking.steps_since_precondition = 0;
                    }
                    if tracking.precondition_tick.is_some() {
                        if goal.eval(ctx) {
                            tracking.precondition_tick = None;
                            tracking.steps_since_precondition = 0;
                        } else {
                            tracking.steps_since_precondition += 1;
                            if tracking.steps_since_precondition > *bound {
                                *violated = true;
                                self.violations.push(PropertyViolation {
                                    property_name: prop.name(),
                                    tick: ctx.tick,
                                    details: format!(
                                        "goal not reached within {bound} steps after precondition"
                                    ),
                                });
                            }
                        }
                    }
                }
                _ => {}
            }
        }

        self.last_trace_len = ctx.trace.len();
    }
}

/// Every send is eventually received (within the session).
#[must_use]
pub fn send_recv_liveness(sid: SessionId, bound: usize) -> Property {
    Property::SendRecvLiveness { sid, bound }
}

/// Type state never regresses (monotonic advancement).
#[must_use]
pub fn type_monotonicity(sid: SessionId) -> Property {
    Property::TypeMonotonicity { sid }
}

/// No session faults.
#[must_use]
pub fn no_faults() -> Property {
    Property::NoFaults
}

/// Buffer occupancy bounded.
#[must_use]
pub fn buffer_bound(sid: SessionId, max: usize) -> Property {
    Property::BufferBound { sid, max }
}

fn check_no_faults(ctx: &PropertyContext<'_>) -> bool {
    ctx.coroutines
        .iter()
        .all(|c| !matches!(c.status, CoroStatus::Faulted(_)))
}

fn check_simplex(ctx: &PropertyContext<'_>) -> bool {
    let one = FixedQ32::one();
    let tol = tolerance();
    let neg_tol = -tol;
    for coro in ctx.coroutines {
        let vals = registers_to_f64s(&coro.regs);
        if vals.is_empty() {
            continue;
        }
        let sum: FixedQ32 = vals.iter().sum();
        if (sum - one).abs() > tol {
            return false;
        }
        if vals.iter().any(|v| *v < neg_tol) {
            return false;
        }
    }
    true
}

fn distance_to_equilibrium(ctx: &PropertyContext<'_>) -> Option<FixedQ32> {
    let mut max_dist: Option<FixedQ32> = None;
    for coro in ctx.coroutines {
        let vals = registers_to_f64s(&coro.regs);
        if vals.is_empty() {
            continue;
        }
        let n = FixedQ32::try_from_usize(vals.len()).unwrap_or(FixedQ32::one());
        let target = FixedQ32::one() / n;
        let dist = vals
            .iter()
            .map(|v| (*v - target).powi(2))
            .sum::<FixedQ32>()
            .sqrt();
        max_dist = Some(max_dist.map_or(dist, |m| m.max(dist)));
    }
    max_dist
}

// RECURSION_SAFE: structural recursion over a finite local type tree.
fn contains_recursion(lt: &LocalTypeR) -> bool {
    match lt {
        LocalTypeR::Mu { body, .. } => contains_recursion(body),
        LocalTypeR::Var(_) => true,
        LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
            branches.iter().any(|(_, _, cont)| contains_recursion(cont))
        }
        LocalTypeR::End => false,
    }
}

// RECURSION_SAFE: structural recursion over a finite local type tree.
fn type_depth(lt: &LocalTypeR) -> usize {
    match lt {
        LocalTypeR::End | LocalTypeR::Var(_) => 0,
        LocalTypeR::Mu { body, .. } => type_depth(body),
        LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
            let max_branch = branches
                .iter()
                .map(|(_, _, cont)| type_depth(cont))
                .max()
                .unwrap_or(0);
            1 + max_branch
        }
    }
}

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

    #[test]
    fn parse_property_handles_parameterized_invariants() {
        match parse_property("buffer_bound(3,7)").expect("parse property") {
            Property::BufferBound { sid, max } => {
                assert_eq!(sid, 3);
                assert_eq!(max, 7);
            }
            other => panic!("expected buffer bound, got {other:?}"),
        }

        match parse_property("type_monotonicity(5)").expect("parse property") {
            Property::TypeMonotonicity { sid } => assert_eq!(sid, 5),
            other => panic!("expected type monotonicity, got {other:?}"),
        }
    }
}
