// Copyright 2026 @quantus/codec authors & contributors
// SPDX-License-Identifier: Apache-2.0

//! Addressing chain state, using the runtime's own description of where it lives.
//!
//! A storage key is `twox128(pallet) ‖ twox128(item)`, then each map key hashed
//! by the hasher the entry declares. None of those choices are known here:
//! the pallet prefix, the item name, the hashers, the key type and the value
//! type all come out of the metadata, so a runtime upgrade that re-hashes a map
//! or changes a value's shape is followed rather than mis-read.

use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec::Vec;

use frame_metadata::v14::{StorageEntryType, StorageHasher};

use crate::runtime::{CodecError, Runtime};

/// Where a storage value lives and what it decodes as.
#[derive(Debug, Clone)]
pub struct StorageTarget {
    /// The full key, ready for `state_getStorage`.
    pub key: Vec<u8>,
    /// The registry type its value decodes as.
    pub value_ty: u32,
    /// What the chain means when it returns nothing.
    ///
    /// `Some(bytes)` for a `Default` entry — an unfunded account reads as a zero
    /// balance, not as an error. `None` for an `Optional` entry, where nothing
    /// means nothing. Conflating the two is how a wallet reports "failed to load"
    /// for an account that simply has no money in it.
    pub default: Option<Vec<u8>>,
}

impl Runtime {
    /// Resolve a storage entry, hashing any map keys as the runtime declares.
    ///
    /// `keys` are JSON, interpreted against the key types the metadata gives —
    /// so an `AccountId32` is the hex string its inner array accepts, and a
    /// double map takes two values in the order the entry lists its hashers.
    pub fn storage_target(
        &self,
        pallet: &str,
        item: &str,
        keys: &[serde_json::Value],
    ) -> Result<StorageTarget, CodecError> {
        let entry = self
            .metadata
            .pallets
            .iter()
            .find(|p| p.name == pallet)
            .and_then(|p| p.storage.as_ref())
            .and_then(|s| s.entries.iter().find(|e| e.name == item))
            .ok_or_else(|| CodecError::NoStorageEntry(format!("{pallet}::{item}")))?;

        let prefix = self
            .metadata
            .pallets
            .iter()
            .find(|p| p.name == pallet)
            .and_then(|p| p.storage.as_ref())
            .map(|s| s.prefix.clone())
            .unwrap_or_else(|| pallet.to_string());

        let mut key = twox_128(prefix.as_bytes()).to_vec();

        key.extend_from_slice(&twox_128(item.as_bytes()));

        let value_ty = match &entry.ty {
            StorageEntryType::Plain(ty) => {
                if !keys.is_empty() {
                    return Err(CodecError::NoStorageEntry(format!(
                        "{pallet}::{item} takes no keys"
                    )));
                }

                ty.id
            }
            StorageEntryType::Map {
                hashers,
                key: key_ty,
                value,
            } => {
                if hashers.len() != keys.len() {
                    return Err(CodecError::NoStorageEntry(format!(
                        "{pallet}::{item} takes {} key(s), {} given",
                        hashers.len(),
                        keys.len()
                    )));
                }

                // One hasher means the declared key type *is* the key. More than
                // one means it is a tuple, one element per hasher, and the
                // elements are hashed separately rather than as a unit.
                let key_tys: Vec<u32> = if hashers.len() == 1 {
                    alloc::vec![key_ty.id]
                } else {
                    match self.types().resolve(key_ty.id).map(|t| &t.type_def) {
                        Some(scale_info::TypeDef::Tuple(t)) => {
                            t.fields.iter().map(|f| f.id).collect()
                        }
                        _ => {
                            return Err(CodecError::NoStorageEntry(format!(
                                "{pallet}::{item} has {} hashers and a non-tuple key",
                                hashers.len()
                            )))
                        }
                    }
                };

                for ((supplied, ty), hasher) in keys.iter().zip(key_tys).zip(hashers.iter()) {
                    let value = self.json_to_value(supplied, ty)?;
                    let mut encoded = Vec::new();

                    scale_value::scale::encode_as_type(&value, ty, self.types(), &mut encoded)
                        .map_err(|e| CodecError::Encode(format!("{pallet}::{item} key: {e}")))?;

                    key.extend_from_slice(&hash_key(hasher, &encoded));
                }

                value.id
            }
        };

        let default = match &entry.modifier {
            frame_metadata::v14::StorageEntryModifier::Default => Some(entry.default.clone()),
            frame_metadata::v14::StorageEntryModifier::Optional => None,
        };

        Ok(StorageTarget {
            default,
            key,
            value_ty,
        })
    }

    /// Decode a storage value against the type its entry declares.
    ///
    /// `bytes` is what `state_getStorage` returned, or the entry's default when
    /// it returned nothing.
    pub fn decode_storage_value(
        &self,
        value_ty: u32,
        bytes: &[u8],
    ) -> Result<serde_json::Value, CodecError> {
        let mut cursor = bytes;
        let value = self
            .decode_checked(value_ty, &mut cursor)
            .map_err(|e| CodecError::Decode(e))?;

        if !cursor.is_empty() {
            return Err(CodecError::Decode(format!(
                "{} trailing bytes after storage value",
                cursor.len()
            )));
        }

        Ok(self.render(&value))
    }
}

/// `twox128`, as Substrate uses it for pallet and item prefixes.
fn twox_128(input: &[u8]) -> [u8; 16] {
    use twox_hash::XxHash64;

    let mut out = [0u8; 16];

    out[..8].copy_from_slice(&XxHash64::oneshot(0, input).to_le_bytes());
    out[8..].copy_from_slice(&XxHash64::oneshot(1, input).to_le_bytes());

    out
}

/// Hash one map key the way its entry declares.
///
/// The `Concat` variants keep the key after its hash, which is what makes a map
/// enumerable. The plain variants do not.
fn hash_key(hasher: &StorageHasher, encoded: &[u8]) -> Vec<u8> {
    use blake2::digest::consts::{U16, U32};
    use blake2::{Blake2b, Digest};
    use twox_hash::XxHash64;

    match hasher {
        StorageHasher::Blake2_128 => Blake2b::<U16>::digest(encoded).to_vec(),
        StorageHasher::Blake2_256 => Blake2b::<U32>::digest(encoded).to_vec(),
        StorageHasher::Blake2_128Concat => {
            let mut v = Blake2b::<U16>::digest(encoded).to_vec();

            v.extend_from_slice(encoded);
            v
        }
        StorageHasher::Twox128 => twox_128(encoded).to_vec(),
        StorageHasher::Twox256 => {
            let mut v = Vec::with_capacity(32);

            for seed in 0..4u64 {
                v.extend_from_slice(&XxHash64::oneshot(seed, encoded).to_le_bytes());
            }

            v
        }
        StorageHasher::Twox64Concat => {
            let mut v = XxHash64::oneshot(0, encoded).to_le_bytes().to_vec();

            v.extend_from_slice(encoded);
            v
        }
        StorageHasher::Identity => encoded.to_vec(),
    }
}

/// Hex, for the JSON boundary.
pub(crate) fn hex(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(2 + bytes.len() * 2);

    s.push_str("0x");

    for b in bytes {
        s.push(char::from_digit((b >> 4) as u32, 16).expect("nibble"));
        s.push(char::from_digit((b & 0x0f) as u32, 16).expect("nibble"));
    }

    s
}