//! Implementation of the `secrets` module.
//!
//! Provides cryptographically strong helpers backed by OS randomness:
//! - `token_bytes(nbytes=None)` returns random bytes (default 32)
//! - `token_hex(nbytes=None)` returns a lowercase hex token string
//! - `token_urlsafe(nbytes=None)` returns a URL-safe base64 token string without padding
//! - `randbelow(exclusive_upper_bound)` returns a random int in `[0, n)`
//! - `choice(sequence)` returns a random element from a non-empty sequence
//! - `compare_digest(a, b)` performs a timing-safe comparison for supported input kinds
//! - `SystemRandom` is re-exported from `random.SystemRandom`

use rand::RngCore;

use crate::{
    args::ArgValues,
    exception_private::{ExcType, RunResult, SimpleException},
    heap::{DropWithHeap, Heap, HeapData, HeapId},
    intern::{Interns, StaticStrings},
    modules::{ModuleFunctions, random_mod},
    resource::ResourceTracker,
    types::{AttrCallResult, Bytes, Module, PyTrait, Str},
    value::Value,
};

/// Default number of random bytes generated by token helpers.
const DEFAULT_TOKEN_BYTES: i64 = 32;

/// URL-safe base64 alphabet (`+`/`/` replaced by `-`/`_`).
const URLSAFE_BASE64_ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

/// `secrets` module callable exports.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum::Display, serde::Serialize, serde::Deserialize)]
#[strum(serialize_all = "snake_case")]
pub(crate) enum SecretsFunctions {
    TokenBytes,
    TokenHex,
    TokenUrlsafe,
    Randbelow,
    Choice,
    CompareDigest,
}

/// Creates the `secrets` module and registers all supported attributes.
pub fn create_module(
    heap: &mut Heap<impl ResourceTracker>,
    interns: &Interns,
) -> Result<HeapId, crate::resource::ResourceError> {
    let mut module = Module::new(StaticStrings::Secrets);

    for (name, function) in [
        ("token_bytes", SecretsFunctions::TokenBytes),
        ("token_hex", SecretsFunctions::TokenHex),
        ("token_urlsafe", SecretsFunctions::TokenUrlsafe),
        ("randbelow", SecretsFunctions::Randbelow),
        ("choice", SecretsFunctions::Choice),
        ("compare_digest", SecretsFunctions::CompareDigest),
    ] {
        module.set_attr_str(
            name,
            Value::ModuleFunction(ModuleFunctions::Secrets(function)),
            heap,
            interns,
        )?;
    }

    let system_random = clone_system_random_from_random_module(heap, interns)?;
    module.set_attr_str("SystemRandom", system_random, heap, interns)?;

    heap.allocate(HeapData::Module(module))
}

/// Dispatches `secrets` module function calls.
pub(super) fn call(
    heap: &mut Heap<impl ResourceTracker>,
    interns: &Interns,
    function: SecretsFunctions,
    args: ArgValues,
) -> RunResult<AttrCallResult> {
    match function {
        SecretsFunctions::TokenBytes => token_bytes(heap, interns, args),
        SecretsFunctions::TokenHex => token_hex(heap, interns, args),
        SecretsFunctions::TokenUrlsafe => token_urlsafe(heap, interns, args),
        SecretsFunctions::Randbelow => randbelow(heap, args),
        SecretsFunctions::Choice => choice(heap, interns, args),
        SecretsFunctions::CompareDigest => compare_digest(heap, interns, args),
    }
}

/// Implements `secrets.token_bytes(nbytes=None)`.
fn token_bytes(heap: &mut Heap<impl ResourceTracker>, interns: &Interns, args: ArgValues) -> RunResult<AttrCallResult> {
    let size = parse_optional_nbytes_arg(args, "secrets.token_bytes", heap, interns)?;
    let mut bytes = vec![0u8; size];
    rand::rngs::OsRng.fill_bytes(&mut bytes);
    let bytes_id = heap.allocate(HeapData::Bytes(Bytes::new(bytes)))?;
    Ok(AttrCallResult::Value(Value::Ref(bytes_id)))
}

/// Implements `secrets.token_hex(nbytes=None)`.
fn token_hex(heap: &mut Heap<impl ResourceTracker>, interns: &Interns, args: ArgValues) -> RunResult<AttrCallResult> {
    let size = parse_optional_nbytes_arg(args, "secrets.token_hex", heap, interns)?;
    let mut bytes = vec![0u8; size];
    rand::rngs::OsRng.fill_bytes(&mut bytes);
    let token = encode_hex_lower(&bytes);
    let token_id = heap.allocate(HeapData::Str(Str::from(token)))?;
    Ok(AttrCallResult::Value(Value::Ref(token_id)))
}

/// Implements `secrets.token_urlsafe(nbytes=None)`.
fn token_urlsafe(
    heap: &mut Heap<impl ResourceTracker>,
    interns: &Interns,
    args: ArgValues,
) -> RunResult<AttrCallResult> {
    let size = parse_optional_nbytes_arg(args, "secrets.token_urlsafe", heap, interns)?;
    let mut bytes = vec![0u8; size];
    rand::rngs::OsRng.fill_bytes(&mut bytes);
    let token = encode_urlsafe_base64_no_padding(&bytes);
    let token_id = heap.allocate(HeapData::Str(Str::from(token)))?;
    Ok(AttrCallResult::Value(Value::Ref(token_id)))
}

/// Implements `secrets.randbelow(exclusive_upper_bound)`.
fn randbelow(heap: &mut Heap<impl ResourceTracker>, args: ArgValues) -> RunResult<AttrCallResult> {
    let upper_bound_value = args.get_one_arg("secrets.randbelow", heap)?;
    let upper_bound = int_like_to_i64(&upper_bound_value, heap);
    upper_bound_value.drop_with_heap(heap);
    let upper_bound = upper_bound?;

    if upper_bound <= 0 {
        return Err(SimpleException::new_msg(ExcType::ValueError, "Upper bound must be positive.").into());
    }

    let upper_bound_u64 = u64::try_from(upper_bound).expect("validated positive i64 fits u64");
    let value_u64 = os_randbelow_u64(upper_bound_u64);
    let value_i64 = i64::try_from(value_u64).expect("randbelow result is strictly below i64 upper bound");
    Ok(AttrCallResult::Value(Value::Int(value_i64)))
}

/// Implements `secrets.choice(sequence)`.
fn choice(heap: &mut Heap<impl ResourceTracker>, interns: &Interns, args: ArgValues) -> RunResult<AttrCallResult> {
    let mut sequence = args.get_one_arg("secrets.choice", heap)?;

    // Match CPython's `SystemRandom.choice` shape: length check first, then index.
    let Some(length) = sequence.py_len(heap, interns) else {
        let type_name = sequence.py_type(heap);
        sequence.drop_with_heap(heap);
        return Err(ExcType::type_error(format!(
            "object of type '{type_name}' has no len()"
        )));
    };
    if length == 0 {
        sequence.drop_with_heap(heap);
        return Err(SimpleException::new_msg(ExcType::IndexError, "Cannot choose from an empty sequence").into());
    }

    let index_u64 = os_randbelow_u64(length as u64);
    let index_i64 = i64::try_from(index_u64).expect("usize length is expected to fit i64 in this runtime");
    let item_result = sequence.py_getitem(&Value::Int(index_i64), heap, interns);
    sequence.drop_with_heap(heap);
    let item = item_result?;

    Ok(AttrCallResult::Value(item))
}

/// Implements `secrets.compare_digest(a, b)`.
fn compare_digest(
    heap: &mut Heap<impl ResourceTracker>,
    interns: &Interns,
    args: ArgValues,
) -> RunResult<AttrCallResult> {
    let (left, right) = args.get_two_args("secrets.compare_digest", heap)?;
    let result = compare_digest_impl(&left, &right, heap, interns);
    left.drop_with_heap(heap);
    right.drop_with_heap(heap);
    result.map(|value| AttrCallResult::Value(Value::Bool(value)))
}

/// Parses `nbytes` from `token_*` functions, supporting positional and keyword forms.
fn parse_optional_nbytes_arg(
    args: ArgValues,
    function_name: &str,
    heap: &mut Heap<impl ResourceTracker>,
    interns: &Interns,
) -> RunResult<usize> {
    let (mut positional, kwargs) = args.into_parts();
    let positional_count = positional.len();
    if positional_count > 1 {
        positional.drop_with_heap(heap);
        kwargs.drop_with_heap(heap);
        return Err(ExcType::type_error_at_most(function_name, 1, positional_count));
    }
    let positional_nbytes = positional.next();

    let mut keyword_nbytes = None;
    for (key, value) in kwargs {
        let Some(key_name) = key.as_either_str(heap).map(|s| s.as_str(interns).to_owned()) else {
            key.drop_with_heap(heap);
            value.drop_with_heap(heap);
            positional_nbytes.drop_with_heap(heap);
            keyword_nbytes.drop_with_heap(heap);
            return Err(ExcType::type_error_kwargs_nonstring_key());
        };
        key.drop_with_heap(heap);

        if key_name != "nbytes" {
            value.drop_with_heap(heap);
            positional_nbytes.drop_with_heap(heap);
            keyword_nbytes.drop_with_heap(heap);
            return Err(ExcType::type_error_unexpected_keyword(function_name, &key_name));
        }
        if keyword_nbytes.is_some() {
            value.drop_with_heap(heap);
            positional_nbytes.drop_with_heap(heap);
            keyword_nbytes.drop_with_heap(heap);
            return Err(ExcType::type_error_multiple_values(function_name, "nbytes"));
        }
        keyword_nbytes = Some(value);
    }

    if positional_nbytes.is_some() && keyword_nbytes.is_some() {
        positional_nbytes.drop_with_heap(heap);
        keyword_nbytes.drop_with_heap(heap);
        return Err(ExcType::type_error_duplicate_arg(function_name, "nbytes"));
    }

    let nbytes_value = positional_nbytes.or(keyword_nbytes);
    let nbytes = parse_nbytes_value(nbytes_value.as_ref(), heap);
    nbytes_value.drop_with_heap(heap);
    nbytes
}

/// Converts the optional `nbytes` value to a validated output length.
fn parse_nbytes_value(nbytes_value: Option<&Value>, heap: &Heap<impl ResourceTracker>) -> RunResult<usize> {
    let nbytes = match nbytes_value {
        None => DEFAULT_TOKEN_BYTES,
        Some(Value::None) => DEFAULT_TOKEN_BYTES,
        Some(value) => int_like_to_i64(value, heap)?,
    };
    if nbytes < 0 {
        return Err(SimpleException::new_msg(ExcType::ValueError, "negative argument not allowed").into());
    }
    Ok(usize::try_from(nbytes).expect("validated non-negative i64 length fits usize on this platform"))
}

/// Converts a Python integer-like value to `i64`, accepting booleans as 0/1.
fn int_like_to_i64(value: &Value, heap: &Heap<impl ResourceTracker>) -> RunResult<i64> {
    match value {
        Value::Bool(b) => Ok(i64::from(*b)),
        _ => value.as_int(heap),
    }
}

/// Clones `random.SystemRandom` for exposure as `secrets.SystemRandom`.
fn clone_system_random_from_random_module(
    heap: &mut Heap<impl ResourceTracker>,
    interns: &Interns,
) -> Result<Value, crate::resource::ResourceError> {
    let random_module_id = random_mod::create_module(heap, interns)?;
    let system_random = {
        let HeapData::Module(random_module) = heap.get(random_module_id) else {
            unreachable!("random_mod::create_module must return a module object");
        };
        random_module
            .attrs()
            .get_by_str("SystemRandom", heap, interns)
            .expect("random module must expose SystemRandom")
            .clone_with_heap(heap)
    };
    Value::Ref(random_module_id).drop_with_heap(heap);
    Ok(system_random)
}

/// Extracts a string view for `str`/interned-string values.
fn value_as_str<'a>(value: &'a Value, heap: &'a Heap<impl ResourceTracker>, interns: &'a Interns) -> Option<&'a str> {
    match value {
        Value::InternString(id) => Some(interns.get_str(*id)),
        Value::Ref(id) => match heap.get(*id) {
            HeapData::Str(s) => Some(s.as_str()),
            _ => None,
        },
        _ => None,
    }
}

/// Extracts bytes for bytes-like values supported by the runtime.
fn value_as_bytes(value: &Value, heap: &Heap<impl ResourceTracker>, interns: &Interns) -> Option<Vec<u8>> {
    match value {
        Value::InternBytes(id) => Some(interns.get_bytes(*id).to_vec()),
        Value::Ref(id) => match heap.get(*id) {
            HeapData::Bytes(bytes) => Some(bytes.as_slice().to_vec()),
            _ => None,
        },
        _ => None,
    }
}

/// Implements compare_digest type handling and timing-safe comparison behavior.
fn compare_digest_impl(
    left: &Value,
    right: &Value,
    heap: &Heap<impl ResourceTracker>,
    interns: &Interns,
) -> RunResult<bool> {
    let left_str = value_as_str(left, heap, interns);
    let right_str = value_as_str(right, heap, interns);

    if let (Some(left_str), Some(right_str)) = (left_str, right_str) {
        if !left_str.is_ascii() || !right_str.is_ascii() {
            return Err(ExcType::type_error(
                "comparing strings with non-ASCII characters is not supported",
            ));
        }
        return Ok(constant_time_eq(left_str.as_bytes(), right_str.as_bytes()));
    }

    let left_bytes = value_as_bytes(left, heap, interns);
    let right_bytes = value_as_bytes(right, heap, interns);

    if let (Some(left_bytes), Some(right_bytes)) = (left_bytes.as_ref(), right_bytes.as_ref()) {
        return Ok(constant_time_eq(left_bytes, right_bytes));
    }

    let left_is_str = left_str.is_some();
    let right_is_str = right_str.is_some();
    let left_is_bytes = left_bytes.is_some();
    let right_is_bytes = right_bytes.is_some();

    if (left_is_str && right_is_bytes) || (left_is_bytes && right_is_str) {
        return Err(ExcType::type_error("a bytes-like object is required, not 'str'"));
    }

    if left_is_str || right_is_str {
        return Err(ExcType::type_error(format!(
            "unsupported operand types(s) or combination of types: '{}' and '{}'",
            left.py_type(heap),
            right.py_type(heap),
        )));
    }

    if left_is_bytes {
        return Err(ExcType::type_error(format!(
            "a bytes-like object is required, not '{}'",
            right.py_type(heap),
        )));
    }
    if right_is_bytes {
        return Err(ExcType::type_error(format!(
            "a bytes-like object is required, not '{}'",
            left.py_type(heap),
        )));
    }

    Err(ExcType::type_error(format!(
        "unsupported operand types(s) or combination of types: '{}' and '{}'",
        left.py_type(heap),
        right.py_type(heap),
    )))
}

/// Performs a timing-safe equality check over two byte sequences.
fn constant_time_eq(left: &[u8], right: &[u8]) -> bool {
    let max_len = left.len().max(right.len());
    let mut diff = left.len() ^ right.len();
    for index in 0..max_len {
        let left_byte = left.get(index).copied().unwrap_or(0);
        let right_byte = right.get(index).copied().unwrap_or(0);
        diff |= usize::from(left_byte ^ right_byte);
    }
    diff == 0
}

/// Encodes bytes as lowercase hexadecimal text.
fn encode_hex_lower(bytes: &[u8]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut out = String::with_capacity(bytes.len() * 2);
    for &byte in bytes {
        out.push(char::from(HEX[(byte >> 4) as usize]));
        out.push(char::from(HEX[(byte & 0x0f) as usize]));
    }
    out
}

/// Encodes bytes into URL-safe base64 and strips any `=` padding.
fn encode_urlsafe_base64_no_padding(bytes: &[u8]) -> String {
    let mut out = String::with_capacity((bytes.len() * 4).div_ceil(3));
    for chunk in bytes.chunks(3) {
        let b0 = chunk[0];
        let b1 = if chunk.len() > 1 { chunk[1] } else { 0 };
        let b2 = if chunk.len() > 2 { chunk[2] } else { 0 };
        let n = (u32::from(b0) << 16) | (u32::from(b1) << 8) | u32::from(b2);

        out.push(char::from(URLSAFE_BASE64_ALPHABET[((n >> 18) & 0x3f) as usize]));
        out.push(char::from(URLSAFE_BASE64_ALPHABET[((n >> 12) & 0x3f) as usize]));
        if chunk.len() > 1 {
            out.push(char::from(URLSAFE_BASE64_ALPHABET[((n >> 6) & 0x3f) as usize]));
        } else {
            out.push('=');
        }
        if chunk.len() > 2 {
            out.push(char::from(URLSAFE_BASE64_ALPHABET[(n & 0x3f) as usize]));
        } else {
            out.push('=');
        }
    }
    while out.ends_with('=') {
        out.pop();
    }
    out
}

/// Returns up to 64 random bits from OS entropy.
fn os_getrandbits_u64(k: u32) -> u64 {
    if k == 0 {
        return 0;
    }
    let mut rng = rand::rngs::OsRng;
    let mut bits_remaining = k;
    let mut shift = 0u32;
    let mut value = 0u64;
    while bits_remaining > 0 {
        let take = bits_remaining.min(32);
        let mut word = rng.next_u32();
        if take < 32 {
            word >>= 32 - take;
        }
        value |= u64::from(word) << shift;
        shift += 32;
        bits_remaining -= take;
    }
    value
}

/// Returns an unbiased integer in `[0, n)` using OS entropy.
fn os_randbelow_u64(n: u64) -> u64 {
    if n <= 1 {
        return 0;
    }
    let k = u64::BITS - n.leading_zeros();
    loop {
        let random = os_getrandbits_u64(k);
        if random < n {
            return random;
        }
    }
}
