//! Multiplication operations for BigInt.
//!
//! Implements constant-time multiplication using Comba method for small operands
//! and Karatsuba algorithm for large operands (threshold n≥16).

use crate::error::{BigIntError, Result};
use crate::limbs::{Limb, limb_add_carry, limb_mac, limb_mul_wide, limb_sub_borrow};
use crate::types::BigIntCore;

#[cfg(feature = "alloc")]
use crate::types::BigInt;
#[cfg(feature = "alloc")]
use alloc::vec;

/// Karatsuba threshold: switch to Karatsuba when limb count >= 16.
const KARATSUBA_THRESHOLD: usize = 16;

/// Comba multiplication: standard schoolbook method.
///
/// Computes c = a × b where a and b have n limbs.
/// Result c has 2n limbs.
fn comba_mul(a: &[Limb], b: &[Limb], c: &mut [Limb]) {
    let n = a.len();
    debug_assert_eq!(b.len(), n);
    debug_assert!(c.len() >= 2 * n);

    // Initialize result to zero
    for i in 0..(2 * n) {
        c[i] = 0;
    }

    // Comba multiplication algorithm
    for i in 0..n {
        let mut carry = 0u64;
        for j in 0..n {
            let wide = limb_mul_wide(a[i], b[j]);
            let (lo, hi) = (wide.lo, wide.hi);
            let (sum_lo, carry1) = limb_mac(c[i + j], lo, carry);
            c[i + j] = sum_lo;
            carry = carry1 + hi;
        }
        c[i + n] = carry;
    }
}

/// Karatsuba multiplication for large operands.
///
/// This algorithm splits operands in half and uses the identity:
/// (a_high * B^m + a_low) * (b_high * B^m + b_low) =
/// a_high*b_high * B^(2m) + (a_high*b_low + a_low*b_high) * B^m + a_low*b_low
///
/// Which can be computed as:
/// z0 = a_low * b_low
/// z2 = a_high * b_high
/// z1 = (a_low + a_high) * (b_low + b_high) - z0 - z2
/// result = z0 + z1 * B^m + z2 * B^(2m)
///
/// For constant-time execution, we ensure all operations have fixed sizes.
fn karatsuba_mul(a: &[Limb], b: &[Limb], c: &mut [Limb]) {
    let n = a.len();

    // Base case: use Comba for small operands
    if n < KARATSUBA_THRESHOLD {
        comba_mul(a, b, c);
        return;
    }

    // Split operands in half (constant-time split)
    let m = n / 2;
    let (a_low, a_high) = a.split_at(m);
    let (b_low, b_high) = b.split_at(m);

    // Compute z0 = a_low * b_low (m limbs)
    let mut z0 = vec![0u64; 2 * m];
    if m >= KARATSUBA_THRESHOLD {
        karatsuba_mul(a_low, b_low, &mut z0);
    } else {
        comba_mul(a_low, b_low, &mut z0);
    }

    // Compute z2 = a_high * b_high (m limbs)
    let mut z2 = vec![0u64; 2 * m];
    if m >= KARATSUBA_THRESHOLD {
        karatsuba_mul(a_high, b_high, &mut z2);
    } else {
        comba_mul(a_high, b_high, &mut z2);
    }

    // Compute a_low + a_high and b_low + b_high (with carry handling)
    let mut a_sum = vec![0u64; m + 1]; // +1 for potential carry
    let mut carry = 0u64;
    for i in 0..m {
        let (sum, carry_out) = limb_add_carry(a_low[i], a_high[i], carry);
        a_sum[i] = sum;
        carry = carry_out;
    }
    if carry > 0 {
        a_sum[m] = carry;
    }

    let mut b_sum = vec![0u64; m + 1]; // +1 for potential carry
    carry = 0;
    for i in 0..m {
        let (sum, carry_out) = limb_add_carry(b_low[i], b_high[i], carry);
        b_sum[i] = sum;
        carry = carry_out;
    }
    if carry > 0 {
        b_sum[m] = carry;
    }

    // Compute z1 = (a_sum) * (b_sum) (up to 2m + 2 limbs)
    let sum_len = m + if carry > 0 { 1 } else { 0 };
    let mut z1 = vec![0u64; 2 * sum_len];
    if sum_len >= KARATSUBA_THRESHOLD {
        // Pad to next power of 2 for Karatsuba
        let padded_len = sum_len.next_power_of_two();
        let mut a_padded = vec![0u64; padded_len];
        let mut b_padded = vec![0u64; padded_len];
        a_padded[..sum_len].copy_from_slice(&a_sum[..sum_len]);
        b_padded[..sum_len].copy_from_slice(&b_sum[..sum_len]);
        let mut z1_padded = vec![0u64; 2 * padded_len];
        karatsuba_mul(&a_padded, &b_padded, &mut z1_padded);
        z1.copy_from_slice(&z1_padded[..2 * sum_len]);
    } else {
        comba_mul(&a_sum[..sum_len], &b_sum[..sum_len], &mut z1[..2 * sum_len]);
    }

    // Subtract z0 and z2 from z1: z1 = z1 - z0 - z2
    let mut borrow = 0u64;
    for i in 0..(2 * m) {
        let z0_val = if i < z0.len() { z0[i] } else { 0 };
        let z2_val = if i < z2.len() { z2[i] } else { 0 };
        let (diff1, borrow1) = limb_sub_borrow(z1[i], z0_val, borrow);
        let (diff2, borrow2) = limb_sub_borrow(diff1, z2_val, borrow1);
        z1[i] = diff2;
        borrow = borrow2;
    }

    // Handle final borrow if any
    for i in (2 * m)..z1.len() {
        if borrow > 0 {
            let (diff, borrow_out) = limb_sub_borrow(z1[i], 0, borrow);
            z1[i] = diff;
            borrow = borrow_out;
        }
    }

    // Combine results: c = z0 + z1 * B^m + z2 * B^(2m)
    // First, copy z0 to result
    for i in 0..(2 * m) {
        if i < c.len() {
            c[i] = z0[i];
        }
    }

    // Add z1 * B^m (shifted by m limbs)
    carry = 0;
    for i in 0..z1.len() {
        let target_idx = i + m;
        if target_idx < c.len() {
            let (sum, carry_out) = limb_add_carry(c[target_idx], z1[i], carry);
            c[target_idx] = sum;
            carry = carry_out;
        }
    }

    // Add z2 * B^(2m) (shifted by 2*m limbs)
    carry = 0;
    for i in 0..z2.len() {
        let target_idx = i + 2 * m;
        if target_idx < c.len() {
            let (sum, carry_out) = limb_add_carry(c[target_idx], z2[i], carry);
            c[target_idx] = sum;
            carry = carry_out;
        }
    }

    // Handle any remaining carry
    let mut idx = 2 * m + z2.len();
    while carry > 0 && idx < c.len() {
        let (sum, carry_out) = limb_add_carry(c[idx], 0, carry);
        c[idx] = sum;
        carry = carry_out;
        idx += 1;
    }
}

/// Multiply two BigInt values.
#[cfg(feature = "alloc")]
pub fn mul(a: &BigInt, b: &BigInt) -> Result<BigInt> {
    // Handle zero cases
    if a.is_zero() || b.is_zero() {
        let max_limbs = a.max_limbs().max(b.max_limbs());
        return Ok(BigInt::new(max_limbs));
    }

    let n = a.limb_count().max(b.limb_count());
    let result_limbs = 2 * n;

    // Check for overflow
    if result_limbs > a.max_limbs().max(b.max_limbs()) {
        return Err(BigIntError::Overflow);
    }

    let max_limbs = a.max_limbs().max(b.max_limbs()).max(result_limbs);
    let mut result = BigInt::new(max_limbs);
    result.ensure_capacity(result_limbs)?;

    let a_limbs = a.limbs();
    let b_limbs = b.limbs();

    // Normalize to same length for multiplication
    let mut a_normalized = vec![0u64; n];
    let mut b_normalized = vec![0u64; n];
    a_normalized[..a_limbs.len()].copy_from_slice(a_limbs);
    b_normalized[..b_limbs.len()].copy_from_slice(b_limbs);

    let result_limbs_mut = result.limbs_mut();

    // Choose algorithm based on limb count
    if n >= KARATSUBA_THRESHOLD {
        karatsuba_mul(&a_normalized, &b_normalized, result_limbs_mut);
    } else {
        comba_mul(&a_normalized, &b_normalized, result_limbs_mut);
    }

    // Set sign: sign(c) = sign(a) XOR sign(b)
    result.set_sign(a.sign() ^ b.sign());
    result.canonicalize()?;

    Ok(result)
}

/// Square a BigInt value.
///
/// This is slightly more efficient than general multiplication.
#[cfg(feature = "alloc")]
pub fn sqr(a: &BigInt) -> Result<BigInt> {
    mul(a, a)
}

/// In-place multiplication for dynamic BigInt.
#[cfg(feature = "alloc")]
pub fn mul_assign(a: &mut BigInt, b: &BigInt) -> Result<()> {
    let result = mul(a, b)?;
    *a = result;
    Ok(())
}

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

    #[test]
    fn test_mul_simple() {
        let a = BigInt::from_u64(10, 10);
        let b = BigInt::from_u64(20, 10);
        let result = mul(&a, &b).unwrap();
        assert_eq!(result.limbs()[0], 200);
    }

    #[test]
    fn test_mul_zero() {
        let a = BigInt::from_u64(10, 10);
        let b = BigInt::from_u64(0, 10);
        let result = mul(&a, &b).unwrap();
        assert!(result.is_zero());
    }
}
