use std::borrow::Cow;
use std::io::Write;
use std::mem::ManuallyDrop;
use std::sync::{Mutex, OnceLock, PoisonError};

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::{Error, Extension, Runtime, RuntimeOptions, SnapshotError, extensions, op};

const MAGIC: &[u8; 8] = b"AVRILSNP";
const FORMAT: u32 = 1;
const HEADER: usize = 16;

/// Static extension heaps that share one V8 read-only heap, captured before
/// live initialization or guest execution.
#[derive(Clone)]
pub struct Snapshot {
    bytes: Cow<'static, [u8]>,
    blob: usize,
    manifest: Manifest,
}

#[derive(Clone, Serialize, Deserialize)]
struct Manifest {
    /// SHA-256 of the V8 startup data, identifying its read-only heap.
    heap: [u8; 32],
    layouts: Vec<Layout>,
}

#[derive(Clone, Serialize, Deserialize)]
pub(crate) struct Layout {
    fingerprint: [u8; 32],
    pub(crate) context: usize,
    pub(crate) state: extensions::SnapshotState,
}

impl Snapshot {
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }

    /// Import a trusted snapshot artifact.
    ///
    /// # Safety
    /// The bytes must be an intact artifact from `Runtime::create_snapshot`,
    /// produced with the same avril-core sources, target, V8 build, and flags,
    /// except for the generator-only random seed set by `initialize_snapshot_generator`.
    /// A build-time generator may be a different executable; its native reference
    /// catalog must match exactly. Only the avril container is validated here;
    /// snapshots are not an untrusted input format.
    pub unsafe fn from_bytes(bytes: impl Into<Cow<'static, [u8]>>) -> Result<Self, SnapshotError> {
        let bytes = bytes.into();
        if bytes.len() < HEADER || bytes[..8] != MAGIC[..] {
            return Err(SnapshotError::InvalidArtifact("not an avril snapshot"));
        }
        let field = |at: usize| u32::from_le_bytes(bytes[at..at + 4].try_into().unwrap());
        if field(8) != FORMAT {
            return Err(SnapshotError::InvalidArtifact(
                "unsupported snapshot format",
            ));
        }
        let blob = usize::try_from(field(12))
            .ok()
            .and_then(|length| HEADER.checked_add(length))
            .filter(|blob| *blob <= bytes.len())
            .ok_or(SnapshotError::InvalidArtifact(
                "truncated snapshot manifest",
            ))?;
        let manifest = serde_json::from_slice(&bytes[HEADER..blob])?;
        Ok(Self { bytes, blob, manifest })
    }

    pub(crate) fn layout(&self, fingerprint: &[u8; 32]) -> Result<&Layout, SnapshotError> {
        self.manifest
            .layouts
            .iter()
            .find(|layout| layout.fingerprint == *fingerprint)
            .ok_or(SnapshotError::IncompatibleLayout)
    }

    fn startup_data(&self) -> v8::StartupData {
        match &self.bytes {
            Cow::Borrowed(bytes) => {
                let bytes: &'static [u8] = bytes;
                Cow::Borrowed(&bytes[self.blob..])
            }
            Cow::Owned(bytes) => Cow::Owned(bytes[self.blob..].to_vec()),
        }
        .into()
    }
}

/// The read-only heap a V8 isolate is created from. V8 shares one read-only heap
/// between all live isolates of a process and requires every isolate to bring a
/// byte-identical image; separately generated artifacts never do (see SNAPSHOTS.md).
#[derive(Clone, Copy, PartialEq, Eq)]
enum Heap {
    BuiltIn,
    Snapshot([u8; 32]),
}

static LIVE_HEAP: Mutex<Option<(Heap, usize)>> = Mutex::new(None);

/// Accounts for one isolate on the process read-only heap; dropped after its isolate.
struct HeapLease(Heap);

impl HeapLease {
    fn acquire(heap: Heap) -> Result<Self, SnapshotError> {
        let mut live = LIVE_HEAP.lock().unwrap_or_else(PoisonError::into_inner);
        match &mut *live {
            Some((current, isolates)) if *current == heap => *isolates += 1,
            Some(_) => return Err(SnapshotError::ReadOnlyHeapInUse),
            None => *live = Some((heap, 1)),
        }
        Ok(Self(heap))
    }
}

impl Drop for HeapLease {
    fn drop(&mut self) {
        let mut live = LIVE_HEAP.lock().unwrap_or_else(PoisonError::into_inner);
        if let Some((heap, isolates)) = &mut *live {
            debug_assert!(*heap == self.0);
            *isolates -= 1;
            if *isolates == 0 {
                *live = None;
            }
        }
    }
}

#[derive(Clone, Copy)]
pub(crate) enum Source<'a> {
    BuiltIn,
    Creator,
    Snapshot(&'a Snapshot),
}

// rusty_v8 has no public query for the thread's entered isolate; it calls this
// binding.cc export for its own OwnedIsolate drop check. Keep it this one sealed joint.
unsafe extern "C" {
    fn v8__Isolate__GetCurrent() -> *mut std::ffi::c_void;
}

fn is_current(isolate: &v8::Isolate) -> bool {
    // SAFETY: UnsafeRawIsolatePtr is repr(transparent) over the isolate pointer, the
    // same `v8::Isolate*` that GetCurrent returns; only addresses are compared.
    let isolate = unsafe {
        std::mem::transmute::<v8::UnsafeRawIsolatePtr, *mut std::ffi::c_void>(
            isolate.as_raw_isolate_ptr(),
        )
    };
    // SAFETY: GetCurrent only reads V8's thread-local entered isolate.
    unsafe { v8__Isolate__GetCurrent() == isolate }
}

// Snapshot creators cannot be dropped without create_blob in rusty_v8. Keep that
// ownership exception here, rather than spreading ManuallyDrop through Runtime.
pub(crate) struct IsolateOwner {
    isolate: Option<v8::OwnedIsolate>,
    // Released only once V8 has deleted the isolate: explicitly in Drop, or after
    // create_blob by field order.
    lease: Option<HeapLease>,
    creator: bool,
}

impl IsolateOwner {
    /// Rejects an incompatible read-only heap before V8 creates the isolate.
    pub(crate) fn new(source: Source) -> Result<Self, SnapshotError> {
        let lease = HeapLease::acquire(match source {
            Source::BuiltIn | Source::Creator => Heap::BuiltIn,
            Source::Snapshot(snapshot) => Heap::Snapshot(snapshot.manifest.heap),
        })?;
        let references = external_references().into();
        let mut isolate = match source {
            Source::BuiltIn => {
                v8::Isolate::new(v8::Isolate::create_params().external_references(references))
            }
            Source::Creator => v8::Isolate::snapshot_creator(Some(references), None),
            Source::Snapshot(snapshot) => v8::Isolate::new(
                v8::Isolate::create_params()
                    .external_references(references)
                    .snapshot_blob(snapshot.startup_data()),
            ),
        };
        if !matches!(source, Source::Snapshot(_)) {
            crate::microtasks::add_listener(&mut isolate);
        }
        Ok(Self {
            isolate: Some(isolate),
            lease: Some(lease),
            creator: matches!(source, Source::Creator),
        })
    }

    pub(crate) fn is_creator(&self) -> bool {
        self.creator
    }

    pub(crate) fn release_native_state(&mut self) {
        release_native_state(self);
    }

    fn create_blob(mut self) -> Option<v8::StartupData> {
        assert!(self.creator);
        self.release_native_state();
        self.isolate
            .take()
            .unwrap()
            .create_blob(v8::FunctionCodeHandling::Keep)
    }
}

impl std::ops::Deref for IsolateOwner {
    type Target = v8::OwnedIsolate;

    fn deref(&self) -> &Self::Target {
        self.isolate.as_ref().unwrap()
    }
}

impl std::ops::DerefMut for IsolateOwner {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.isolate.as_mut().unwrap()
    }
}

impl Drop for IsolateOwner {
    fn drop(&mut self) {
        // After create_blob, V8 has deleted the creator and the lease field releases normally.
        let Some(isolate) = self.isolate.take() else {
            return;
        };
        // A panicking disposal, such as rusty_v8's out-of-order drop check, leaves the
        // native isolate holding V8's read-only heap; keep it accounted until exit.
        let lease = ManuallyDrop::new(self.lease.take());
        if !self.creator {
            drop(isolate);
            drop(ManuallyDrop::into_inner(lease));
            return;
        }
        // Like Deno, abandon the creator rather than panic or serialize a failed heap.
        // It keeps V8's built-in read-only heap alive, so its lease stays pinned too.
        let mut isolate = ManuallyDrop::new(isolate);
        // Native state may own isolates entered after the creator, so release it first,
        // but still leave the creator if one of its destructors panics.
        let released = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            release_native_state(&mut isolate)
        }));
        if is_current(&isolate) {
            // V8 entered the creator on construction and only exits it in the destructor
            // we must skip. Restore the thread's previous isolate for runtimes created earlier.
            // SAFETY: The creator is this thread's current isolate, and it is never used again.
            unsafe { isolate.exit() };
        }
        let _ = writeln!(
            std::io::stderr(),
            "WARNING: abandoned snapshot isolate was leaked"
        );
        // A second panic during unwinding would abort; the first one is already propagating.
        if let Err(panic) = released
            && !std::thread::panicking()
        {
            std::panic::resume_unwind(panic);
        }
    }
}

/// Drops core-owned native state, including every global handle it holds.
fn release_native_state(isolate: &mut v8::Isolate) {
    crate::tasks::clear(isolate);
    op::clear(isolate);
    crate::microtasks::clear(isolate);
    crate::reporting::clear(isolate);
    crate::stack_trace::clear(isolate);
    crate::modules::clear(isolate);
    extensions::clear(isolate);
    isolate.remove_slot::<std::rc::Rc<std::cell::RefCell<crate::WasmStreamingRegistry>>>();
    isolate.remove_slot::<op::RegisteredOps>();
}

/// Identity of every layout, including the core-owned modules and native catalog
/// shared by all of them. Fields are length-prefixed under a domain tag.
pub(crate) fn fingerprint(extensions: &[Extension]) -> Result<[u8; 32], Error> {
    let order = extensions::validate(extensions)?;
    let mut hash = Sha256::new();
    field(&mut hash, b"avril-core snapshot layout v1");
    hash.update(core_identity());
    let entries: Vec<_> = order
        .iter()
        .filter_map(|index| extensions[*index].entry_point)
        .collect();
    count(&mut hash, entries.len());
    for entry in entries {
        field(&mut hash, entry.as_bytes());
    }
    let mut sources: Vec<_> = extensions
        .iter()
        .flat_map(|extension| extension.modules.iter().copied())
        .collect();
    sources.sort_unstable_by_key(|(name, _)| *name);
    count(&mut hash, sources.len());
    for (name, source) in sources {
        field(&mut hash, name.as_bytes());
        field(&mut hash, source.as_bytes());
    }
    let mut ops: Vec<_> = extensions
        .iter()
        .flat_map(|extension| extension.ops.iter().map(|op| op.name))
        .collect();
    ops.sort_unstable();
    count(&mut hash, ops.len());
    for name in ops {
        field(&mut hash, name.as_bytes());
    }
    Ok(hash.finalize().into())
}

fn core_identity() -> &'static [u8; 32] {
    static IDENTITY: OnceLock<[u8; 32]> = OnceLock::new();
    IDENTITY.get_or_init(|| {
        let mut hash = Sha256::new();
        field(&mut hash, b"avril-core");
        field(&mut hash, include_str!("primordials.js").as_bytes());
        count(&mut hash, extensions::SOURCES.len());
        for (name, source) in extensions::SOURCES {
            field(&mut hash, name.as_bytes());
            field(&mut hash, source.as_bytes());
        }
        let modules = extensions::native_modules();
        count(&mut hash, modules.len());
        for (name, exports) in modules {
            field(&mut hash, name.as_bytes());
            count(&mut hash, exports.len());
            for (export, _) in exports {
                field(&mut hash, export.as_bytes());
            }
        }
        for (group, references) in external_reference_groups() {
            field(&mut hash, group.as_bytes());
            count(&mut hash, references.len());
        }
        hash.finalize().into()
    })
}

fn count(hash: &mut Sha256, count: usize) {
    hash.update((count as u64).to_le_bytes());
}

fn field(hash: &mut Sha256, bytes: &[u8]) {
    count(hash, bytes.len());
    hash.update(bytes);
}

fn external_references() -> Vec<v8::ExternalReference> {
    external_reference_groups()
        .into_iter()
        .flat_map(|(_, references)| references)
        .collect()
}

fn external_reference_groups() -> [(&'static str, Vec<v8::ExternalReference>); 4] {
    [
        ("ops", op::RegisteredOps::external_references()),
        ("extensions", extensions::external_references()),
        ("modules", crate::modules::external_references()),
        ("microtasks", crate::microtasks::external_references()),
    ]
}

pub(crate) fn create(layouts: Vec<RuntimeOptions>) -> Result<Snapshot, Error> {
    let fingerprints = layouts
        .iter()
        .map(|options| fingerprint(&options.extensions))
        .collect::<Result<Vec<_>, _>>()?;
    for (index, fingerprint) in fingerprints.iter().enumerate() {
        if fingerprints[..index].contains(fingerprint) {
            return Err(SnapshotError::DuplicateLayout.into());
        }
    }
    let mut isolate = IsolateOwner::new(Source::Creator)?;
    {
        v8::scope!(let scope, &mut *isolate);
        let pristine = v8::Context::new(scope, Default::default());
        scope.set_default_context(pristine);
    }
    let mut captured = Vec::with_capacity(layouts.len());
    for (options, fingerprint) in layouts.into_iter().zip(fingerprints) {
        let runtime = Runtime::bootstrap(isolate, options, crate::ContextSource::Fresh)?;
        let (owner, context, state) = capture(runtime)?;
        isolate = owner;
        captured.push(Layout { fingerprint, context, state });
    }
    let blob = isolate.create_blob().ok_or(SnapshotError::CreationFailed)?;
    let manifest = serde_json::to_vec(&Manifest {
        heap: Sha256::digest(&*blob).into(),
        layouts: captured,
    })
    .map_err(SnapshotError::Metadata)?;
    let length = u32::try_from(manifest.len())
        .map_err(|_| SnapshotError::InvalidArtifact("snapshot manifest is too large"))?;
    let mut bytes = Vec::with_capacity(HEADER + manifest.len() + blob.len());
    bytes.extend_from_slice(MAGIC);
    bytes.extend_from_slice(&FORMAT.to_le_bytes());
    bytes.extend_from_slice(&length.to_le_bytes());
    bytes.extend_from_slice(&manifest);
    bytes.extend_from_slice(&blob);
    // SAFETY: These bytes were just assembled from this creator's blob.
    Ok(unsafe { Snapshot::from_bytes(bytes) }?)
}

/// Adds one static layout as a snapshot context, then tears down its native state
/// so the creator can bootstrap the next layout independently.
fn capture(
    mut runtime: Runtime,
) -> Result<(IsolateOwner, usize, extensions::SnapshotState), Error> {
    if runtime.has_pending_work() {
        return Err(SnapshotError::ActiveState("pending work").into());
    }
    if crate::modules::has_modules(&runtime.isolate) {
        return Err(SnapshotError::ActiveState("guest modules").into());
    }
    if crate::reporting::has_handlers(&runtime.isolate) {
        return Err(SnapshotError::ActiveState("event handlers").into());
    }
    if !runtime
        .isolate
        .get_slot::<std::rc::Rc<std::cell::RefCell<crate::WasmStreamingRegistry>>>()
        .unwrap()
        .borrow()
        .is_empty()
    {
        return Err(SnapshotError::ActiveState("Wasm streaming").into());
    }
    let (context, state) = runtime.with_context(|scope| {
        if let Err(error) = remove_stack_trace_limit(scope)? {
            return Some(Err(error));
        }
        let state = extensions::capture(scope);
        let context = scope.get_current_context();
        Some(Ok((scope.add_context(context), state)))
    })??;
    let Runtime { context: handle, mut isolate, .. } = runtime;
    drop(handle);
    isolate.release_native_state();
    Ok((isolate, context, state))
}

/// V8 installs `Error.stackTraceLimit` only in live contexts, so static bootstrap
/// gets V8's default to give failures locations.
pub(crate) fn define_stack_trace_limit(scope: &mut v8::PinScope) -> Option<()> {
    let (error, key) = stack_trace_limit(scope)?;
    let limit = v8::Integer::new(scope, STACK_TRACE_LIMIT);
    error.create_data_property(scope, key.into(), limit.into())?;
    Some(())
}

/// Restored contexts receive V8's own live property, so static code must leave the
/// default untouched for them to match cold bootstraps.
fn remove_stack_trace_limit(scope: &mut v8::PinScope) -> Option<Result<(), SnapshotError>> {
    let (error, key) = stack_trace_limit(scope)?;
    let limit = error.get(scope, key.into())?;
    let unchanged = limit.is_int32() && limit.int32_value(scope)? == STACK_TRACE_LIMIT;
    if !unchanged || !error.delete(scope, key.into())? {
        return Some(Err(SnapshotError::ActiveState("Error.stackTraceLimit")));
    }
    Some(Ok(()))
}

/// V8's default `--stack-trace-limit`.
const STACK_TRACE_LIMIT: i32 = 10;

fn stack_trace_limit<'s>(
    scope: &mut v8::PinScope<'s, '_>,
) -> Option<(v8::Local<'s, v8::Object>, v8::Local<'s, v8::String>)> {
    let global = scope.get_current_context().global(scope);
    let key = v8::String::new(scope, "Error")?;
    let error = v8::Local::<v8::Object>::try_from(global.get(scope, key.into())?).ok()?;
    Some((error, v8::String::new(scope, "stackTraceLimit")?))
}