use alloc::{format, vec::Vec};
#[cfg(any(test, feature = "testing"))]
use core::ops::Range;

use miden_air::{
    MidenMultiAir, ProverStatement, PublicInputs, StarkConfig, Statement, config, debug,
    trace::{MainTrace, decoder::NUM_USER_OP_HELPERS},
};
use miden_core::{
    deferred::{Digest, TRUE_DIGEST},
    program::ExecutionClaim,
    serde::{
        BudgetedReader, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
        SliceReader,
    },
};

use crate::{
    Felt, MIN_STACK_DEPTH, ProgramInfo, StackInputs, StackOutputs, Word, ZERO,
    fast::ExecutionOutput, field::QuadFelt, utils::RowMajorMatrix,
};

pub(crate) mod utils;
use utils::ChipletTraceFragment;

pub mod chiplets;
pub(crate) mod execution_tracer;

mod block_stack;
mod parallel;
mod range;
mod stack;
mod trace_state;

#[cfg(test)]
mod tests;

// RE-EXPORTS
// ================================================================================================

pub(crate) use execution_tracer::TraceReplay;
pub use miden_air::trace::RowIndex;
pub use miden_core::deferred::PrecompileWitness;
pub use parallel::{
    CORE_TRACE_WIDTH, DEFAULT_MAX_PROVER_MEMORY_BYTES, build_trace, build_trace_with_budget,
};
// Re-exported for the streaming trace-build path
// (`FastProcessor::execute_and_build_trace_sync`), which is std-only; the buffered path
// uses `build_hasher_chiplet` and `MAX_TRACE_LEN` directly within `parallel`.
#[cfg(feature = "std")]
pub(crate) use parallel::{MAX_TRACE_LEN, build_hasher_chiplet, build_trace_with_prebuilt_hasher};
#[cfg(feature = "std")]
pub(crate) use trace_state::ResolvedHasherOp;
pub use utils::{ChipletsLengths, TraceLenSummary};

/// Complete in-memory witness produced by a traced program execution.
///
/// The processor constructs its VM witness and optional singleton precompile witness from the same
/// execution output, so they retain the same deferred root. The aggregate may contain private and
/// potentially large prover data. Its binary form is trusted replay data: sparse MAST node and
/// digest maps inside the trace replay are not checked against a source `MastForest` commitment;
/// see <https://github.com/0xMiden/miden-vm/issues/3303>.
#[derive(Debug)]
pub struct ExecutionWitness {
    vm: VmWitness,
    precompile: Option<PrecompileWitness>,
}

impl ExecutionWitness {
    pub(crate) fn from_execution(
        program_info: ProgramInfo,
        stack_inputs: StackInputs,
        execution_output: ExecutionOutput,
        trace: TraceReplay,
    ) -> Self {
        let ExecutionOutput {
            stack: stack_outputs,
            advice: _,
            memory: _,
            precompile_witness: precompile,
        } = execution_output;
        let precompile_root =
            precompile.as_ref().map_or(TRUE_DIGEST, PrecompileWitness::root_unchecked);
        let vm = VmWitness {
            program_info,
            stack_inputs,
            stack_outputs,
            trace,
            precompile_root,
        };

        Self { vm, precompile }
    }

    /// Returns the public claim associated with this witness.
    pub fn claim(&self) -> ExecutionClaim {
        self.vm.claim()
    }

    /// Returns whether this witness contains precompile work.
    pub const fn has_precompiles(&self) -> bool {
        self.precompile.is_some()
    }

    /// Consumes this witness and returns its supported low-level proving components.
    ///
    /// The [`VmWitness`] can be passed to [`build_trace`]. The optional [`PrecompileWitness`] is
    /// present only when execution authenticated deferred precompile work.
    pub fn into_parts(self) -> (VmWitness, Option<PrecompileWitness>) {
        (self.vm, self.precompile)
    }

    /// Decodes one execution witness from a potentially adversarial byte slice.
    ///
    /// The reader applies an input-proportional limit to byte consumption and collection
    /// preallocation, then rejects trailing bytes. These checks establish safe transport syntax.
    /// They do not prove that the sparse MAST replay is a subset of a particular source forest
    /// because the witness wire does not carry such a proof.
    ///
    /// Use [`Self::read_from_bytes_trusted`] for bytes retained inside a trusted prover system.
    #[track_caller]
    pub fn read_from_bytes(bytes: &[u8]) -> Result<Self, DeserializationError> {
        let budget = bytes.len().saturating_mul(EXECUTION_WITNESS_BYTE_READ_BUDGET_MULTIPLIER);
        let mut reader = BudgetedReader::new(SliceReader::new(bytes), budget);
        let witness = <Self as Deserializable>::read_from(&mut reader)?;

        if reader.has_more_bytes() {
            return Err(DeserializationError::InvalidValue(
                "extra bytes after execution witness payload".into(),
            ));
        }
        Ok(witness)
    }

    /// Decodes execution witness bytes retained inside a trusted prover system.
    ///
    /// This preserves the original permissive byte-slice behavior. It trusts sparse MAST replay
    /// hashes, accepts trailing bytes, and uses a replay-sized byte budget. Use
    /// [`Self::read_from_bytes`] for bytes received across a trust boundary.
    #[track_caller]
    pub fn read_from_bytes_trusted(bytes: &[u8]) -> Result<Self, DeserializationError> {
        let budget = bytes.len().saturating_mul(EXECUTION_WITNESS_BYTE_READ_BUDGET_MULTIPLIER);
        let mut reader = BudgetedReader::new(SliceReader::new(bytes), budget);
        <Self as Deserializable>::read_from(&mut reader)
    }
}

/// Covers nested replay length checks while keeping untrusted work proportional to input size.
const EXECUTION_WITNESS_BYTE_READ_BUDGET_MULTIPLIER: usize = 4;

/// Current wire format version for [`ExecutionWitness`] serialization.
///
/// The version is written as the first byte of every serialized witness. Deserialization only
/// accepts this exact value; older formats are intentionally unsupported.
const EXECUTION_WITNESS_WIRE_VERSION: u8 = 2;

impl Serializable for ExecutionWitness {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        EXECUTION_WITNESS_WIRE_VERSION.write_into(target);
        self.vm.write_into(target);
        match &self.precompile {
            Some(precompile) => {
                target.write_u8(1);
                precompile.write_into(target);
            },
            None => target.write_u8(0),
        }
    }
}

impl Deserializable for ExecutionWitness {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let version = u8::read_from(source)?;
        if version != EXECUTION_WITNESS_WIRE_VERSION {
            return Err(DeserializationError::InvalidValue(format!(
                "unsupported execution witness wire version {version} (expected \
                 {EXECUTION_WITNESS_WIRE_VERSION})"
            )));
        }
        let vm = VmWitness::read_from(source)?;
        let precompile = match source.read_u8()? {
            0 => {
                if vm.precompile_root != TRUE_DIGEST {
                    return Err(DeserializationError::InvalidValue(
                        "VM witness claims deferred work but no precompile witness is present"
                            .into(),
                    ));
                }
                None
            },
            1 => {
                let witness = PrecompileWitness::read_from(source)?;
                if witness.root_unchecked() != vm.precompile_root {
                    return Err(DeserializationError::InvalidValue(
                        "precompile witness root does not match the VM witness precompile root"
                            .into(),
                    ));
                }
                Some(witness)
            },
            tag => {
                return Err(DeserializationError::InvalidValue(format!(
                    "invalid precompile witness option tag {tag}"
                )));
            },
        };
        Ok(Self { vm, precompile })
    }

    fn read_from_bytes(bytes: &[u8]) -> Result<Self, DeserializationError> {
        ExecutionWitness::read_from_bytes(bytes)
    }
}

/// Witness required to materialize and prove a VM execution trace.
///
/// This potentially large value contains private replay data and is consumed by trace-building and
/// proving operations. Its binary form is trusted replay data: sparse MAST node and digest maps
/// inside the trace replay are not checked against a source `MastForest` commitment; see
/// <https://github.com/0xMiden/miden-vm/issues/3303>.
///
/// Its direct [`Serializable`] encoding is an unversioned implementation detail. Use the
/// versioned [`ExecutionWitness`] envelope for transport or persistence.
#[derive(Debug)]
pub struct VmWitness {
    program_info: ProgramInfo,
    stack_inputs: StackInputs,
    stack_outputs: StackOutputs,
    trace: TraceReplay,
    precompile_root: Digest,
}

impl VmWitness {
    /// Returns the public claim associated with this witness.
    pub fn claim(&self) -> ExecutionClaim {
        ExecutionClaim::from_program_info(
            self.program_info.clone(),
            self.stack_inputs,
            self.stack_outputs,
        )
    }

    /// Returns whether this witness authenticates deferred precompile work.
    pub fn has_precompiles(&self) -> bool {
        self.precompile_root != TRUE_DIGEST
    }

    /// Takes the hasher replay out, leaving an empty buffered one.
    ///
    /// The streaming path uses this to drop the replay's channel sender once execution has
    /// finished, so the concurrently running hasher builder sees its input stream end.
    #[cfg(feature = "std")]
    pub(crate) fn take_hasher_replay(&mut self) -> trace_state::HasherRequestReplay {
        core::mem::take(&mut self.trace.hasher_for_chiplet)
    }

    /// Returns the replay data captured during execution.
    #[cfg(any(test, feature = "testing"))]
    #[cfg_attr(all(feature = "testing", not(test)), expect(dead_code))]
    pub(crate) fn trace_replay(&self) -> &TraceReplay {
        &self.trace
    }

    // Kept for tests that force invalid replay data without widening the public API.
    #[cfg(any(test, feature = "testing"))]
    #[cfg_attr(all(feature = "testing", not(test)), expect(dead_code))]
    pub(crate) fn trace_replay_mut(&mut self) -> &mut TraceReplay {
        &mut self.trace
    }

    /// Returns the number of sparse MAST forests captured in the replay.
    ///
    /// Available under the `testing` feature for external tests that check which replay data a
    /// serialized witness preserves.
    #[cfg(any(test, feature = "testing"))]
    #[allow(dead_code)]
    pub fn mast_forest_count(&self) -> usize {
        self.trace.mast_forest_store.len()
    }
}

impl Serializable for VmWitness {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.program_info.write_into(target);
        self.stack_inputs.write_into(target);
        self.stack_outputs.write_into(target);
        self.trace.write_into(target);
        self.precompile_root.write_into(target);
    }
}

impl Deserializable for VmWitness {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        Ok(Self {
            program_info: ProgramInfo::read_from(source)?,
            stack_inputs: StackInputs::read_from(source)?,
            stack_outputs: StackOutputs::read_from(source)?,
            trace: TraceReplay::read_from(source)?,
            precompile_root: Digest::read_from(source)?,
        })
    }
}

// VM EXECUTION TRACE
// ================================================================================================

/// Execution trace which is generated when a program is executed on the VM.
///
/// The trace consists of the following components:
/// - Per-AIR main traces for Core, Chiplets, Eidos compression, and byte-pair lookup.
/// - Information about the program (program hash and the kernel).
/// - Information about the initial and final stack states and authenticated precompile root.
/// - Summary of trace lengths of the main trace components.
#[derive(Debug)]
pub struct VmTrace {
    main_trace: MainTrace,
    program_info: ProgramInfo,
    stack_inputs: StackInputs,
    stack_outputs: StackOutputs,
    precompile_root: Digest,
    trace_len_summary: TraceLenSummary,
}

impl VmTrace {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------

    pub(crate) fn new_from_parts(
        program_info: ProgramInfo,
        stack_inputs: StackInputs,
        stack_outputs: StackOutputs,
        precompile_root: Digest,
        main_trace: MainTrace,
        trace_len_summary: TraceLenSummary,
    ) -> Self {
        Self {
            main_trace,
            program_info,
            stack_inputs,
            stack_outputs,
            precompile_root,
            trace_len_summary,
        }
    }

    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns the program info of this execution trace.
    pub fn program_info(&self) -> &ProgramInfo {
        &self.program_info
    }

    /// Returns hash of the program execution of which resulted in this execution trace.
    pub fn program_hash(&self) -> &Word {
        self.program_info.program_hash()
    }

    /// Returns outputs of the program execution which resulted in this execution trace.
    pub fn stack_outputs(&self) -> &StackOutputs {
        &self.stack_outputs
    }

    /// Returns the public inputs for this execution trace.
    pub fn public_inputs(&self) -> PublicInputs {
        PublicInputs::new(
            self.program_info.clone(),
            self.stack_inputs,
            self.stack_outputs,
            self.precompile_root,
        )
    }

    /// Returns the public values for this execution trace.
    pub fn to_public_values(&self) -> Vec<Felt> {
        self.public_inputs().to_elements()
    }

    /// Returns a reference to the main trace.
    pub fn main_trace(&self) -> &MainTrace {
        &self.main_trace
    }

    /// Returns a mutable reference to the main trace.
    pub fn main_trace_mut(&mut self) -> &mut MainTrace {
        &mut self.main_trace
    }

    /// Returns the authenticated root of the deferred precompile state.
    pub fn precompile_root(&self) -> Digest {
        self.precompile_root
    }

    /// Returns the owned stack outputs required for proof packaging.
    pub fn into_outputs(self) -> StackOutputs {
        self.stack_outputs
    }

    /// Returns the initial state of the top 16 stack registers.
    pub fn init_stack_state(&self) -> StackInputs {
        self.stack_inputs
    }

    /// Returns the final state of the top 16 stack registers.
    pub fn last_stack_state(&self) -> StackOutputs {
        let last_step = RowIndex::from(self.last_step());
        let mut result = [ZERO; MIN_STACK_DEPTH];
        for (i, result) in result.iter_mut().enumerate() {
            *result = self.main_trace.stack_element(i, last_step);
        }
        result.into()
    }

    /// Returns helper registers state at the specified `clk` of the VM
    pub fn get_user_op_helpers_at(&self, clk: u32) -> [Felt; NUM_USER_OP_HELPERS] {
        let mut result = [ZERO; NUM_USER_OP_HELPERS];
        let row = RowIndex::from(clk);
        for (i, result) in result.iter_mut().enumerate() {
            *result = self.main_trace.helper_register(i, row);
        }
        result
    }

    /// Returns the trace length.
    pub fn get_trace_len(&self) -> usize {
        self.main_trace.num_rows()
    }

    /// Returns the length of the trace (number of rows in the main trace).
    pub fn length(&self) -> usize {
        self.get_trace_len()
    }

    /// Returns a summary of the per-AIR trace lengths.
    pub fn trace_len_summary(&self) -> &TraceLenSummary {
        &self.trace_len_summary
    }

    // DEBUG CONSTRAINT CHECKING
    // --------------------------------------------------------------------------------------------

    /// Validates this execution trace against all AIR constraints without generating a STARK
    /// proof.
    ///
    /// This is the recommended way to test trace correctness. It is much faster than full STARK
    /// proving and provides better error diagnostics (panics on the first constraint violation
    /// with the instance and row index).
    ///
    /// # Panics
    ///
    /// Panics if any AIR constraint evaluates to nonzero.
    pub fn check_constraints(&self) {
        let (core_matrix, chiplets_matrix, eidos_compression_matrix, and8_lookup_matrix) =
            self.main_trace.clone_air_matrices();
        let (public_values, kernel_felts) = self.public_inputs().to_air_inputs();

        let statement: Statement<Felt, QuadFelt, MidenMultiAir> =
            Statement::new(MidenMultiAir::new(), public_values, kernel_felts)
                .expect("statement construction failed");
        let prover_statement = ProverStatement::new(
            statement,
            vec![core_matrix, chiplets_matrix, eidos_compression_matrix, and8_lookup_matrix],
        )
        .expect("prover statement construction failed");

        let config = config::eidos_config(config::pcs_params(), config::RELATION_DIGEST);
        debug::check_constraints(&prover_statement, config.challenger());
    }

    /// Clones the trace buffers into the per-AIR matrices consumed by the multi-AIR prover.
    pub fn clone_air_matrices(
        &self,
    ) -> (
        RowMajorMatrix<Felt>,
        RowMajorMatrix<Felt>,
        RowMajorMatrix<Felt>,
        RowMajorMatrix<Felt>,
    ) {
        self.main_trace.clone_air_matrices()
    }

    /// Consuming variant for the proving hot path: moves the row-major buffers.
    pub fn into_air_matrices(
        self,
    ) -> (
        RowMajorMatrix<Felt>,
        RowMajorMatrix<Felt>,
        RowMajorMatrix<Felt>,
        RowMajorMatrix<Felt>,
    ) {
        self.main_trace.into_air_matrices()
    }

    // HELPER METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns the index of the last row in the Core trace.
    fn last_step(&self) -> usize {
        self.main_trace.core_height() - 1
    }

    #[cfg(any(test, feature = "testing"))]
    pub fn get_column_range(&self, range: Range<usize>) -> Vec<Vec<Felt>> {
        self.main_trace.get_column_range(range)
    }
}

#[cfg(test)]
mod wire_tests {
    use miden_assembly::Assembler;
    use miden_core::{
        deferred::TRUE_DIGEST,
        mast::{BasicBlockNodeBuilder, MastForest},
        operations::Operation,
        program::Program,
    };

    use super::{ExecutionWitness, Serializable, build_trace};
    use crate::{DefaultHost, FastProcessor, Felt, StackInputs, mast::MastNodeId};

    fn execution_witness(source: &str) -> ExecutionWitness {
        let program = Assembler::default()
            .assemble_program("program", source)
            .expect("program should compile")
            .unwrap_program();
        let mut host = DefaultHost::default();
        FastProcessor::new(StackInputs::default())
            .execute_for_proving_sync(&program, &mut host)
            .expect("execution should produce a witness")
    }

    fn deferred_witness() -> ExecutionWitness {
        execution_witness("begin log_deferred end")
    }

    fn deferred_witness_bytes() -> alloc::vec::Vec<u8> {
        deferred_witness().to_bytes()
    }

    #[test]
    fn witness_reports_precompile_state() {
        let plain = execution_witness("begin push.1 drop end");
        assert!(!plain.has_precompiles());

        let deferred = execution_witness("begin log_deferred end");
        assert!(deferred.has_precompiles());
    }

    fn aead_stream_witness() -> ExecutionWitness {
        let mut mast_forest = MastForest::new();
        let basic_block_id = BasicBlockNodeBuilder::new(vec![Operation::CryptoStream])
            .add_to_forest(&mut mast_forest)
            .unwrap();
        mast_forest.make_root(basic_block_id);
        let program = Program::new(mast_forest.into(), basic_block_id);

        let stack = [
            1, 2, 3, 4, // K_CTR
            0, // counter
            0, // src_ptr
            8, // dst_ptr
            1, // remaining
            0, 0, 0, 0, 0, 0, 0, 0, // tail
        ]
        .map(Felt::new_unchecked);
        let mut host = DefaultHost::default();
        FastProcessor::new(StackInputs::new(&stack).unwrap())
            .execute_for_proving_sync(&program, &mut host)
            .expect("AEAD stream execution should produce a witness")
    }

    #[test]
    fn witness_wire_rejects_unsupported_version() {
        let mut bytes = deferred_witness_bytes();
        assert!(ExecutionWitness::read_from_bytes(&bytes).is_ok());
        assert_eq!(bytes[0], 2, "Eidos witnesses must use wire version 2");

        // Rejected versions are checked before decoding replay or witness payloads.
        for version in [0, 1, super::EXECUTION_WITNESS_WIRE_VERSION + 1] {
            bytes[0] = version;
            let err = ExecutionWitness::read_from_bytes(&bytes)
                .expect_err("unsupported witness format must be rejected");
            assert!(format!("{err:?}").contains("unsupported execution witness wire version"));
        }
    }

    #[test]
    fn witness_wire_rejects_trailing_bytes() {
        let mut bytes = deferred_witness_bytes();
        bytes.push(0);

        let err = ExecutionWitness::read_from_bytes(&bytes)
            .expect_err("witness payload with trailing bytes should be rejected");
        assert!(
            format!("{err:?}").contains("extra bytes after execution witness payload"),
            "unexpected error: {err:?}"
        );

        assert!(
            ExecutionWitness::read_from_bytes_trusted(&bytes).is_ok(),
            "the explicit trusted reader should preserve the old permissive behavior"
        );
    }

    #[test]
    fn witness_wire_accepts_large_minimally_encoded_continuation_stack() {
        let mut witness = deferred_witness();
        let continuation = &mut witness
            .vm
            .trace_replay_mut()
            .core_trace_contexts
            .first_mut()
            .expect("witness should contain a trace fragment")
            .continuation;
        for _ in 0..4096 {
            continuation.push_start_node(MastNodeId::from(0));
        }

        let bytes = witness.to_bytes();
        let restored = ExecutionWitness::read_from_bytes(&bytes)
            .expect("valid witness should fit its input-proportional allocation budget");

        assert_eq!(restored.to_bytes(), bytes);
    }

    #[test]
    fn aead_stream_witness_round_trip_rebuilds_the_same_trace() {
        let witness = aead_stream_witness();
        let bytes = witness.to_bytes();
        let claim = witness.claim();
        let (original_vm, original_precompile) = witness.into_parts();
        assert!(original_precompile.is_none());
        let original_trace = build_trace(original_vm).expect("original trace should build");

        let restored =
            ExecutionWitness::read_from_bytes(&bytes).expect("witness should round trip");
        assert_eq!(restored.claim(), claim);
        let (restored_vm, restored_precompile) = restored.into_parts();
        assert!(restored_precompile.is_none());
        let restored_trace = build_trace(restored_vm).expect("restored trace should build");

        assert_eq!(restored_trace.public_inputs(), original_trace.public_inputs());
        assert_eq!(restored_trace.trace_len_summary(), original_trace.trace_len_summary());
        assert_eq!(restored_trace.clone_air_matrices(), original_trace.clone_air_matrices());
    }

    #[test]
    fn witness_wire_rejects_mismatched_precompile_root() {
        let bytes = deferred_witness_bytes();
        let restored = ExecutionWitness::read_from_bytes(&bytes).expect("witness round trip");
        let (vm, precompile) = restored.into_parts();
        let precompile = precompile.expect("deferred execution should carry a precompile witness");
        assert_ne!(vm.precompile_root, TRUE_DIGEST);

        // Tamper only the VM-side precompile root and re-serialize: the two halves of the wire
        // no longer describe the same execution, so deserialization must reject them.
        let tampered = ExecutionWitness {
            vm: super::VmWitness { precompile_root: TRUE_DIGEST, ..vm },
            precompile: Some(precompile),
        };
        let err = ExecutionWitness::read_from_bytes(&tampered.to_bytes())
            .expect_err("tampered witness should be rejected");
        assert!(
            format!("{err:?}")
                .contains("precompile witness root does not match the VM witness precompile root"),
            "unexpected error: {err:?}"
        );

        // Sanity: the untampered wire still round-trips.
        assert!(ExecutionWitness::read_from_bytes(&bytes).is_ok());
    }

    #[test]
    fn witness_wire_rejects_missing_precompile_witness() {
        let bytes = deferred_witness_bytes();
        let restored = ExecutionWitness::read_from_bytes(&bytes).expect("witness round trip");
        let (vm, precompile) = restored.into_parts();
        assert!(precompile.is_some(), "deferred execution should carry a precompile witness");

        // Drop only the precompile witness while the VM side still claims deferred work: the
        // wire must not validate as a complete execution.
        let stripped = ExecutionWitness { vm, precompile: None };
        let err = ExecutionWitness::read_from_bytes(&stripped.to_bytes())
            .expect_err("witness without its precompile half should be rejected");
        assert!(
            format!("{err:?}")
                .contains("VM witness claims deferred work but no precompile witness is present"),
            "unexpected error: {err:?}"
        );
    }
}