//! Core BigInt types and traits.
//!
//! This module provides both dynamic-length and fixed-length (const-generic)
//! BigInt types for different use cases.

use crate::error::{BigIntError, Result};
use crate::limbs::{Limb, canonicalize, is_zero};

/// Core trait for BigInt operations.
///
/// This trait defines common operations that both dynamic and fixed-length
/// BigInt types must implement.
pub trait BigIntCore {
    /// Get the number of limbs.
    fn limb_count(&self) -> usize;

    /// Get a reference to the limbs.
    fn limbs(&self) -> &[Limb];

    /// Get a mutable reference to the limbs.
    fn limbs_mut(&mut self) -> &mut [Limb];

    /// Get the sign bit (false = positive, true = negative).
    fn sign(&self) -> bool;

    /// Set the sign bit.
    fn set_sign(&mut self, sign: bool);

    /// Check if the value is zero.
    fn is_zero(&self) -> bool {
        is_zero(self.limbs())
    }

    /// Ensure canonical form.
    ///
    /// For dynamic BigInt, this removes leading zeros.
    /// For fixed BigInt, this is a no-op (canonicalization not required).
    fn canonicalize(&mut self) -> Result<()>;
}

#[cfg(feature = "alloc")]
/// Dynamic-length BigInt.
///
/// This type uses heap allocation and can grow up to `max_limbs`.
/// It enforces canonical form (no leading zeros, zero sign rules).
#[derive(Clone, Debug)]
pub struct BigInt {
    /// Sign bit: false = positive, true = negative.
    sign: bool,
    /// Limb array in little-endian order.
    limbs: alloc::vec::Vec<Limb>,
    /// Maximum number of limbs allowed.
    max_limbs: usize,
}

#[cfg(feature = "alloc")]
impl BigInt {
    /// Create a new BigInt with the specified maximum limb count.
    pub fn new(max_limbs: usize) -> Self {
        Self {
            sign: false,
            limbs: alloc::vec![0],
            max_limbs,
        }
    }

    /// Create a BigInt from a u64 value.
    pub fn from_u64(value: u64, max_limbs: usize) -> Self {
        if value == 0 {
            Self {
                sign: false,
                limbs: alloc::vec![0],
                max_limbs,
            }
        } else {
            Self {
                sign: false,
                limbs: alloc::vec![value],
                max_limbs,
            }
        }
    }

    /// Create a BigInt from a slice of limbs.
    ///
    /// The limbs are assumed to be in little-endian order.
    pub fn from_limbs(limbs: &[Limb], max_limbs: usize) -> Result<Self> {
        if limbs.is_empty() {
            return Ok(Self {
                sign: false,
                limbs: alloc::vec![0],
                max_limbs,
            });
        }

        if limbs.len() > max_limbs {
            return Err(BigIntError::Overflow);
        }

        let mut limbs_vec = limbs.to_vec();
        let len = canonicalize(&mut limbs_vec);
        limbs_vec.truncate(len);

        // Check for zero and enforce zero sign rule
        let is_zero = is_zero(&limbs_vec[..len]);
        let sign = if is_zero { false } else { false }; // Start with positive, will be set by caller if needed

        Ok(Self {
            sign,
            limbs: limbs_vec,
            max_limbs,
        })
    }

    /// Get the maximum number of limbs.
    pub fn max_limbs(&self) -> usize {
        self.max_limbs
    }

    /// Ensure capacity for at least `n` limbs.
    pub fn ensure_capacity(&mut self, n: usize) -> Result<()> {
        if n > self.max_limbs {
            return Err(BigIntError::Overflow);
        }
        if self.limbs.len() < n {
            self.limbs.resize(n, 0);
        }
        Ok(())
    }
}

#[cfg(feature = "alloc")]
impl BigIntCore for BigInt {
    fn limb_count(&self) -> usize {
        self.limbs.len()
    }

    fn limbs(&self) -> &[Limb] {
        &self.limbs
    }

    fn limbs_mut(&mut self) -> &mut [Limb] {
        &mut self.limbs
    }

    fn sign(&self) -> bool {
        self.sign
    }

    fn set_sign(&mut self, sign: bool) {
        // Enforce zero sign rule: zero must have sign = false
        if self.is_zero() {
            self.sign = false;
        } else {
            self.sign = sign;
        }
    }

    fn canonicalize(&mut self) -> Result<()> {
        let len = canonicalize(&mut self.limbs);
        self.limbs.truncate(len);

        // Enforce zero sign rule
        if self.is_zero() {
            self.sign = false;
        }

        Ok(())
    }
}

/// Fixed-length BigInt with const-generic limb count.
///
/// This type uses a fixed-size array and is suitable for VM registers
/// and fixed-width cryptographic types.
#[derive(Clone, Copy, Debug)]
pub struct BigIntFixed<const L: usize> {
    /// Sign bit: false = positive, true = negative.
    sign: bool,
    /// Fixed-size limb array in little-endian order.
    limbs: [Limb; L],
}

impl<const L: usize> BigIntFixed<L> {
    /// Create a new zero BigIntFixed.
    pub const fn zero() -> Self {
        Self {
            sign: false,
            limbs: [0; L],
        }
    }

    /// Create a BigIntFixed from a u64 value.
    pub fn from_u64(value: u64) -> Self {
        let mut limbs = [0u64; L];
        if value != 0 {
            limbs[0] = value;
        }
        Self { sign: false, limbs }
    }

    /// Create a BigIntFixed from a slice of limbs.
    ///
    /// If the slice is longer than L, this will truncate.
    /// If shorter, the remaining limbs will be zero.
    pub fn from_limbs(limbs: &[Limb]) -> Result<Self> {
        let mut result_limbs = [0u64; L];
        let copy_len = limbs.len().min(L);
        result_limbs[..copy_len].copy_from_slice(&limbs[..copy_len]);

        // Check for overflow if input was longer
        if limbs.len() > L {
            // Check if any truncated limbs were non-zero
            if limbs[L..].iter().any(|&x| x != 0) {
                return Err(BigIntError::Overflow);
            }
        }

        Ok(Self {
            sign: false,
            limbs: result_limbs,
        })
    }

    /// Get a reference to the limb array.
    pub fn as_limbs(&self) -> &[Limb; L] {
        &self.limbs
    }

    /// Get a mutable reference to the limb array.
    pub fn as_limbs_mut(&mut self) -> &mut [Limb; L] {
        &mut self.limbs
    }
}

impl<const L: usize> BigIntCore for BigIntFixed<L> {
    fn limb_count(&self) -> usize {
        L
    }

    fn limbs(&self) -> &[Limb] {
        &self.limbs
    }

    fn limbs_mut(&mut self) -> &mut [Limb] {
        &mut self.limbs
    }

    fn sign(&self) -> bool {
        self.sign
    }

    fn set_sign(&mut self, sign: bool) {
        // Enforce zero sign rule
        if self.is_zero() {
            self.sign = false;
        } else {
            self.sign = sign;
        }
    }

    fn canonicalize(&mut self) -> Result<()> {
        // Fixed-length BigInt skips canonicalization per spec
        // But we still enforce zero sign rule
        if self.is_zero() {
            self.sign = false;
        }
        Ok(())
    }
}

/// Type alias for 256-bit BigInt (4 limbs).
pub type U256 = BigIntFixed<4>;

/// Type alias for 512-bit BigInt (8 limbs).
pub type U512 = BigIntFixed<8>;

/// Type alias for 1024-bit BigInt (16 limbs).
pub type U1024 = BigIntFixed<16>;

/// Type alias for 2048-bit BigInt (32 limbs).
pub type U2048 = BigIntFixed<32>;

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

    #[test]
    fn test_bigint_from_u64() {
        let bi = BigInt::from_u64(42, 10);
        assert_eq!(bi.limbs(), &[42]);
        assert!(!bi.sign());
        assert!(!bi.is_zero());
    }

    #[test]
    fn test_bigint_zero() {
        let bi = BigInt::from_u64(0, 10);
        assert_eq!(bi.limbs(), &[0]);
        assert!(!bi.sign());
        assert!(bi.is_zero());
    }

    #[test]
    fn test_bigint_fixed_from_u64() {
        let bi = U256::from_u64(42);
        assert_eq!(bi.limbs()[0], 42);
        assert!(!bi.sign());
    }

    #[test]
    fn test_bigint_fixed_zero() {
        let bi = U256::zero();
        assert!(bi.is_zero());
        assert!(!bi.sign());
    }

    #[test]
    fn test_canonicalize_zero_sign() {
        let mut bi = BigInt::from_u64(0, 10);
        bi.set_sign(true); // Try to set negative
        bi.canonicalize().unwrap();
        assert!(!bi.sign()); // Should be forced to positive
    }
}
