//! Iterative constraint propagation solver for bit-width inference.
//!
//! Propagates width constraints to a fixpoint using monotonically increasing
//! updates. Bounded by MAX_PROPAGATION_ROUNDS * constraint_count iterations.
//!
//! Termination proof: widths only increase (never decrease), maximum value
//! is 64, so total increases across all rounds <= 64 * node_count.

#![forbid(unsafe_code)]

use super::constraint::WidthConstraint;
use super::types::{FlatNode, Width, WidthDiag, MAX_FLAT_NODES};

/// Maximum rounds of constraint propagation before declaring convergence failure.
/// With widths bounded to [0, 64] and monotonic increase, 16 rounds is generous.
const MAX_PROPAGATION_ROUNDS: usize = 16;

/// Result of the constraint solving pass.
pub struct SolveResult {
    /// Resolved width for each flat node (indexed by NodeId).
    pub widths: Vec<Width>,
    /// Number of propagation rounds executed.
    pub rounds: usize,
    /// Diagnostics emitted during solving (width overflow, unresolved nodes).
    pub diagnostics: Vec<WidthDiag>,
}

/// Seed the initial width array from literals and fixed constraints.
///
/// Returns a width array of length `node_count`, initialized to Width(0)
/// (unresolved) for all nodes, then overwritten by Fixed constraints.
///
/// Bounded: iterates once over constraints.
fn seed_widths(node_count: usize, constraints: &[WidthConstraint]) -> Vec<u32> {
    let mut widths = vec![0u32; node_count];

    for c in constraints {
        if let WidthConstraint::Fixed { node, width } = c {
            let idx = *node as usize;
            if idx < widths.len() {
                widths[idx] = *width;
            }
        }
        if let WidthConstraint::Boolean { node } = c {
            let idx = *node as usize;
            if idx < widths.len() {
                widths[idx] = 1;
            }
        }
    }

    widths
}

/// Run iterative constraint propagation to fixpoint.
///
/// `nodes` and `constraints` must have been generated by `flatten_expr` and
/// `generate_constraints` respectively.
///
/// Bounded: outer loop runs at most MAX_PROPAGATION_ROUNDS times.
/// Inner loop iterates once over constraints per round.
pub fn solve(nodes: &[FlatNode], constraints: &[WidthConstraint]) -> SolveResult {
    let node_count = nodes.len().min(MAX_FLAT_NODES);
    let mut widths = seed_widths(node_count, constraints);
    let mut diagnostics: Vec<WidthDiag> = Vec::new();
    let mut rounds = 0usize;

    for _round in 0..MAX_PROPAGATION_ROUNDS {
        rounds += 1;
        let mut changed = false;

        for c in constraints {
            let new_w = evaluate_constraint(c, &widths);
            if let Some((idx, w)) = new_w {
                if idx < widths.len() {
                    // Standard constraints (Add, Mul, etc.) only increase width.
                    // Narrowed constraints (BitwiseAnd) can decrease width.
                    let is_narrowing = matches!(c, WidthConstraint::Narrowed { .. });
                    if is_narrowing {
                        if w != widths[idx] {
                            widths[idx] = w;
                            changed = true;
                        }
                    } else if w > widths[idx] {
                        widths[idx] = w;
                        changed = true;
                    }
                }
            }
        }

        if !changed {
            break; // Fixpoint reached.
        }
    }

    // Post-solve validation: check for unresolved nodes and width overflow.
    validate_widths(&widths, nodes, &mut diagnostics);

    let resolved: Vec<Width> = widths.iter().map(|w| Width(*w)).collect();

    SolveResult { widths: resolved, rounds, diagnostics }
}

/// Evaluate a single constraint, returning the (node_index, proposed_width)
/// if the constraint wants to update a node, or `None` if no update.
fn evaluate_constraint(c: &WidthConstraint, widths: &[u32]) -> Option<(usize, u32)> {
    match c {
        WidthConstraint::Fixed { node, width } => Some((*node as usize, *width)),
        WidthConstraint::Boolean { node } => Some((*node as usize, 1)),
        WidthConstraint::MaxPlusOne { node, left, right } => {
            let lw = get_w(widths, *left);
            let rw = get_w(widths, *right);
            let w = lw.max(rw).saturating_add(1);
            Some((*node as usize, w))
        }
        WidthConstraint::MaxOf { node, left, right } => {
            let lw = get_w(widths, *left);
            let rw = get_w(widths, *right);
            Some((*node as usize, lw.max(rw)))
        }
        WidthConstraint::SumOf { node, left, right } => {
            let lw = get_w(widths, *left);
            let rw = get_w(widths, *right);
            let w = lw.saturating_add(rw);
            Some((*node as usize, w))
        }
        WidthConstraint::SumAll { node, elements } => {
            let mut total = 0u32;
            for elem in elements {
                total = total.saturating_add(get_w(widths, *elem));
            }
            Some((*node as usize, total))
        }
        WidthConstraint::LeftPlusConst { node, left, shift_amount } => {
            let lw = get_w(widths, *left);
            let w = lw.saturating_add(*shift_amount);
            Some((*node as usize, w))
        }
        WidthConstraint::LeftPlusMaxShift { node, left } => {
            let lw = get_w(widths, *left);
            let w = lw.saturating_add(63);
            Some((*node as usize, w))
        }
        WidthConstraint::LeftMinusConst { node, left, shift_amount } => {
            // Right shift by k narrows the result to max(1, left_width - k).
            // Monotone in left_width: as lw increases, max(1, lw-k) never decreases.
            let lw = get_w(widths, *left);
            let w = lw.saturating_sub(*shift_amount).max(1);
            Some((*node as usize, w))
        }
        WidthConstraint::SameAs { node, source } => {
            let sw = get_w(widths, *source);
            Some((*node as usize, sw))
        }
        WidthConstraint::SameAsPlusOne { node, source } => {
            let sw = get_w(widths, *source);
            Some((*node as usize, sw.saturating_add(1)))
        }
        WidthConstraint::Narrowed { node, source, narrow_width } => {
            let sw = get_w(widths, *source);
            // If source is not yet resolved (0), we can't narrow it yet.
            // But literals and fixed signals are seeded, so this should resolve.
            if sw == 0 {
                Some((*node as usize, *narrow_width))
            } else {
                Some((*node as usize, sw.min(*narrow_width)))
            }
        }
    }
}

/// Safely look up a width by node index.
fn get_w(widths: &[u32], idx: u32) -> u32 {
    widths.get(idx as usize).copied().unwrap_or(0)
}

/// Post-solve validation: check all widths are resolved and within bounds.
///
/// Bounded: iterates once over `widths` (len <= MAX_FLAT_NODES).
fn validate_widths(widths: &[u32], nodes: &[FlatNode], diagnostics: &mut Vec<WidthDiag>) {
    for (i, &w) in widths.iter().enumerate() {
        if i >= MAX_FLAT_NODES {
            break;
        }
        if w == 0 {
            let desc = node_description(i, nodes);
            diagnostics.push(
                WidthDiag::error(format!(
                    "{} node {} ({}) has unresolved width",
                    crate::error_codes::ec(503),
                    i,
                    desc
                ))
                .with_code("E503"),
            );
        }
        if w > Width::MAX.0 {
            let desc = node_description(i, nodes);
            diagnostics.push(
                WidthDiag::error(format!(
                    "{} node {} ({}) requires {} bits, exceeding maximum of 64",
                    crate::error_codes::ec(504),
                    i,
                    desc,
                    w
                ))
                .with_code("E504"),
            );
        }
    }
}

/// Human-readable description of a flat node for diagnostics.
fn node_description(idx: usize, nodes: &[FlatNode]) -> String {
    match nodes.get(idx) {
        Some(FlatNode::Literal { value }) => format!("literal {}", value),
        Some(FlatNode::Signal { name, .. }) => format!("signal '{}'", name),
        Some(FlatNode::Unary { op, .. }) => format!("{:?}", op),
        Some(FlatNode::Binary { op, .. }) => format!("{:?}", op),
        Some(FlatNode::Prev { signal, delay, .. }) => format!("prev('{}', {})", signal, delay),
        Some(FlatNode::UnfoldIndex { name }) => format!("unfold-index '{}'", name),
        _ => "unknown".to_string(),
    }
}

// ---------------------------------------------------------------------------
// Assignment truncation check
// ---------------------------------------------------------------------------

/// Check whether assigning an expression with `expr_width` to a target signal
/// with `target_width` would cause an unsafe truncation.
///
/// `target_signed` indicates whether the target is a signed type,
/// for sign-aware diagnostic messages.
///
/// Returns diagnostics (possibly empty) for this single assignment.
pub fn check_truncation(
    target_name: &str,
    target_width: u32,
    expr_width: Width,
    target_signed: bool,
) -> Vec<WidthDiag> {
    let mut diags = Vec::new();
    if expr_width.0 > target_width {
        let category = if target_signed { "signed" } else { "unsigned" };
        diags.push(
            WidthDiag::error(format!(
                "{} assignment to '{}' truncates {} {} bits to {} bits",
                crate::error_codes::ec(505),
                target_name,
                category,
                expr_width.0,
                target_width
            ))
            .with_code("E505")
            .with_signal(target_name),
        );
    }
    diags
}
