#[cfg(feature = "arbitrary")]
use alloc::vec;
use alloc::{
    borrow::{Cow, ToOwned},
    format,
    string::ToString,
};
use core::{fmt, str::FromStr};

use miden_assembly_syntax::DisplayHex;
use miden_core::serde::{
    ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
};
#[cfg(feature = "arbitrary")]
use proptest::prelude::*;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

/// A unique identifier for optional sections of the Miden package format
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
#[cfg_attr(all(feature = "arbitrary", test), miden_test_serde_macros::serde_test)]
#[repr(transparent)]
pub struct SectionId(Cow<'static, str>);

impl SectionId {
    /// The section containing debug type definitions (primitives, structs, arrays, pointers,
    /// function types)
    pub const DEBUG_TYPES: Self = Self(Cow::Borrowed("debug_types"));
    /// The section containing debug source file paths and checksums
    pub const DEBUG_SOURCES: Self = Self(Cow::Borrowed("debug_sources"));
    /// The section containing debug function metadata
    pub const DEBUG_FUNCTIONS: Self = Self(Cow::Borrowed("debug_functions"));
    /// The section containing source/debug MAST occurrence graph topology
    pub const DEBUG_SOURCE_GRAPH: Self = Self(Cow::Borrowed("debug_source_graph"));
    /// The section containing source-keyed assembly operation and debug variable rows
    pub const DEBUG_SOURCE_MAP: Self = Self(Cow::Borrowed("debug_source_map"));
    /// The section containing assertion error messages keyed by runtime error code
    pub const DEBUG_ERROR_MESSAGES: Self = Self(Cow::Borrowed("debug_error_messages"));
    /// This section provides the encoded metadata for a compiled account component
    ///
    /// Currently, this corresponds to the serialized representation of
    /// `miden-protocol::account::AccountComponentMetadata`, i.e. name, descrioption, storage, that
    /// is associated with this package.
    pub const ACCOUNT_COMPONENT_METADATA: Self = Self(Cow::Borrowed("account_component_metadata"));
    /// This section contains provenance metadata for packages assembled from project sources.
    pub const PROJECT_SOURCE_PROVENANCE: Self = Self(Cow::Borrowed("project_source_provenance"));
    /// This section contains the serialized kernel package linked against by an executable package.
    pub const KERNEL: Self = Self(Cow::Borrowed("kernel"));

    /// Construct a user-defined (i.e. "custom") section identifier
    ///
    /// Section identifiers must be either an ASCII alphanumeric, or one of the following
    /// characters: `.`, `_`, `-`. Additionally, the identifier must start with an ASCII alphabetic
    /// character or `_`.
    pub fn custom(name: impl AsRef<str>) -> Result<Self, InvalidSectionIdError> {
        let name = name.as_ref();
        if !name.starts_with(|c: char| c.is_ascii_alphabetic() || c == '_') {
            return Err(InvalidSectionIdError::InvalidStart);
        }
        if name.contains(|c: char| !c.is_ascii_alphanumeric() && !matches!(c, '.' | '_' | '-')) {
            return Err(InvalidSectionIdError::InvalidCharacter);
        }
        Ok(Self(name.to_string().into()))
    }

    /// Get this section identifier as a string
    #[inline]
    pub fn as_str(&self) -> &str {
        self.0.as_ref()
    }

    /// Returns true if this section contains package debug metadata.
    pub fn is_debug(&self) -> bool {
        matches!(
            self.as_str(),
            "debug_types"
                | "debug_sources"
                | "debug_functions"
                | "debug_source_graph"
                | "debug_source_map"
                | "debug_error_messages"
        )
    }
}

#[derive(Debug, thiserror::Error)]
pub enum InvalidSectionIdError {
    #[error("invalid section id: cannot be empty")]
    Empty,
    #[error(
        "invalid section id: contains invalid characters, only the set [a-z0-9._-] are allowed"
    )]
    InvalidCharacter,
    #[error("invalid section id: must start with a character in the set [a-z_]")]
    InvalidStart,
}

impl FromStr for SectionId {
    type Err = InvalidSectionIdError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "debug_types" => Ok(Self::DEBUG_TYPES),
            "debug_sources" => Ok(Self::DEBUG_SOURCES),
            "debug_functions" => Ok(Self::DEBUG_FUNCTIONS),
            "debug_source_graph" => Ok(Self::DEBUG_SOURCE_GRAPH),
            "debug_source_map" => Ok(Self::DEBUG_SOURCE_MAP),
            "debug_error_messages" => Ok(Self::DEBUG_ERROR_MESSAGES),
            "account_component_metadata" => Ok(Self::ACCOUNT_COMPONENT_METADATA),
            "project_source_provenance" => Ok(Self::PROJECT_SOURCE_PROVENANCE),
            "kernel" => Ok(Self::KERNEL),
            custom => Self::custom(custom),
        }
    }
}

impl fmt::Display for SectionId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

#[derive(Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Section {
    pub id: SectionId,
    pub data: Cow<'static, [u8]>,
}

impl fmt::Debug for Section {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let verbose = f.alternate();
        let mut builder = f.debug_struct("Section");
        builder.field("id", &format_args!("{}", self.id));
        if verbose {
            builder.field("data", &format_args!("{}", DisplayHex(&self.data))).finish()
        } else {
            builder.field("data", &format_args!("{} bytes", self.data.len())).finish()
        }
    }
}

impl Section {
    pub fn new<B>(id: SectionId, data: B) -> Self
    where
        B: Into<Cow<'static, [u8]>>,
    {
        Self { id, data: data.into() }
    }

    /// Returns true if this section is empty, i.e. has no data
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Returns the size in bytes of this section's data
    pub fn len(&self) -> usize {
        self.data.len()
    }
}

impl Serializable for Section {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        let id = self.id.as_str();
        target.write_usize(id.len());
        target.write_bytes(id.as_bytes());
        target.write_usize(self.len());
        target.write_bytes(&self.data);
    }
}

impl Deserializable for Section {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let id_len = source.read_usize()?;
        let id_bytes = source.read_slice(id_len)?;
        let id_str = core::str::from_utf8(id_bytes).map_err(|err| {
            DeserializationError::InvalidValue(format!("invalid utf-8 in section name: {err}"))
        })?;
        let id = SectionId(Cow::Owned(id_str.to_owned()));

        let len = source.read_usize()?;
        let bytes = source.read_slice(len)?;
        Ok(Section { id, data: Cow::Owned(bytes.to_owned()) })
    }
}

#[cfg(feature = "arbitrary")]
impl Arbitrary for SectionId {
    type Parameters = ();
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
        use alloc::string::String;

        let builtins = proptest::sample::select(vec![
            Self::DEBUG_TYPES,
            Self::DEBUG_SOURCES,
            Self::DEBUG_FUNCTIONS,
            Self::DEBUG_SOURCE_GRAPH,
            Self::DEBUG_SOURCE_MAP,
            Self::DEBUG_ERROR_MESSAGES,
            Self::ACCOUNT_COMPONENT_METADATA,
            Self::PROJECT_SOURCE_PROVENANCE,
            Self::KERNEL,
        ]);

        let custom = (
            proptest::prop_oneof![
                proptest::char::range('a', 'z'),
                proptest::char::range('A', 'Z'),
                Just('_'),
            ],
            proptest::collection::vec(
                proptest::prop_oneof![
                    proptest::char::range('a', 'z'),
                    proptest::char::range('A', 'Z'),
                    proptest::char::range('0', '9'),
                    Just('.'),
                    Just('_'),
                    Just('-'),
                ],
                0..31,
            ),
        )
            .prop_map(|(first, rest)| {
                let mut name = String::with_capacity(rest.len() + 1);
                name.push(first);
                name.extend(rest);
                Self::custom(name).expect("generated custom section ids are valid")
            });

        proptest::prop_oneof![builtins, custom].boxed()
    }
}
