alma 0.1.1

use super::{
    resources::{
        WasmtimeUpdateBufferResource, WasmtimeUpdateResources, WasmtimeUpdateViewResource,
        delete_update_resource, resource_handle_error,
    },
    workspace_observe::{
        WasmtimeWorkspaceObserveTaskResourceDrop, WasmtimeWorkspaceObserveTasks,
        workspace_observe_poll_to_guest, workspace_observe_take_to_guest,
    },
};
use crate::plugin::{
    DecodedBufferEdit, DecodedGuestString, DecodedPayload, DecodedWorkspacePath, GuestDecodeField,
    GuestDecodeLimits, PluginBufferHandle, PluginHandleKind, PluginHostImportError,
    PluginHostImportSession, PluginIdentity, PluginOperationalImport, PluginViewHandle,
    PluginWorkspaceObserveTaskCancellationReport, PluginWorkspaceObserveTaskHandle,
    PluginWorkspaceObserveTaskQueueLimit, ValidatedPluginRuntimeLimits, WitHostImport,
    abi::bindings::alma::editor::host::{
        self as wit_host, Host, HostBufferHandle, HostViewHandle, HostWorkspaceObserveTask,
    },
};
use wasmtime::{
    StoreLimits,
    component::{Resource, ResourceTable},
};

/// Wasmtime limit hook state.
pub(super) struct WasmtimeStoreState {
    /// Resource caps.
    pub(super) limits: StoreLimits,
    /// Decode caps for guest-owned import data.
    decode_limits: GuestDecodeLimits,
    /// Host component resources visible during guest calls.
    pub(super) resources: ResourceTable,
    /// Active guest update session and rooted update resources.
    active_update: Option<WasmtimeActiveUpdateState>,
    /// Runtime-owned queue and resources for guest-visible workspace observations.
    pub(super) workspace_observe_tasks: WasmtimeWorkspaceObserveTasks,
    /// First rejected import in the active update.
    import_rejection: Option<WasmtimeHostImportRejection>,
}

impl WasmtimeStoreState {
    /// Creates store state for one component instance.
    pub(super) fn new(limits: StoreLimits, runtime_limits: ValidatedPluginRuntimeLimits) -> Self {
        Self {
            limits,
            decode_limits: GuestDecodeLimits::from_validated_runtime(runtime_limits),
            resources: ResourceTable::new(),
            active_update: None,
            workspace_observe_tasks: WasmtimeWorkspaceObserveTasks::new(
                PluginWorkspaceObserveTaskQueueLimit::default(),
            ),
            import_rejection: None,
        }
    }

    /// Starts one update import session.
    pub(super) fn begin_update(
        &mut self,
        session: PluginHostImportSession,
        view: PluginViewHandle,
        buffer: PluginBufferHandle,
    ) -> Result<WasmtimeUpdateResources, PluginHostImportError> {
        if self.active_update.is_some() {
            return Err(PluginHostImportError::SessionAlreadyActive);
        }
        self.import_rejection = None;
        let _invalidated_stale_resources =
            self.workspace_observe_tasks.invalidate_started_in_update();
        let view = self
            .resources
            .push(view)
            .map(|resource| WasmtimeUpdateViewResource::from_owned(&resource))
            .map_err(|source| resource_handle_error(PluginHandleKind::View.into(), &source))?;
        let buffer = match self.resources.push(buffer) {
            Ok(buffer) => WasmtimeUpdateBufferResource::from_owned(&buffer),
            Err(source) => {
                delete_update_resource(&mut self.resources, view.owned(), PluginHandleKind::View)?;
                return Err(resource_handle_error(
                    PluginHandleKind::Buffer.into(),
                    &source,
                ));
            }
        };
        let resources = WasmtimeUpdateResources::new(view, buffer);
        self.active_update = Some(WasmtimeActiveUpdateState::new(session, resources));
        Ok(resources)
    }

    /// Ends the active update import session.
    pub(super) fn end_update(&mut self) -> Option<PluginHostImportSession> {
        let active_update = self.active_update.take()?;
        let active_update = active_update.into_parts();
        if let Err(error) = active_update.resources.delete_from(&mut self.resources)
            && self.import_rejection.is_none()
        {
            self.import_rejection = Some(WasmtimeHostImportRejection {
                import: PluginOperationalImport::UpdateResourceCleanup,
                source: error,
            });
        }
        Some(active_update.session)
    }

    /// Keeps task resources returned by a successful update live for later polling.
    pub(super) fn commit_workspace_observe_tasks_started_in_update(&mut self) {
        self.workspace_observe_tasks.commit_started_in_update();
    }

    /// Tombstones task resources returned by an update that did not commit.
    pub(super) fn invalidate_workspace_observe_tasks_started_in_update(&mut self) -> usize {
        self.workspace_observe_tasks.invalidate_started_in_update()
    }

    /// Returns the first typed host-import rejection from this update.
    #[cfg(test)]
    pub(super) fn take_import_error(&mut self) -> Option<PluginHostImportError> {
        self.import_rejection
            .take()
            .map(|rejection| rejection.source)
    }

    /// Returns the first typed host-import rejection with adapter context.
    pub(super) const fn take_import_rejection(&mut self) -> Option<WasmtimeHostImportRejection> {
        self.import_rejection.take()
    }

    /// Bounds generated WIT text before it can enter host effects.
    fn status_publish_from_guest(
        &mut self,
        view: &Resource<PluginViewHandle>,
        text: &str,
    ) -> Result<(), PluginHostImportError> {
        let text = DecodedGuestString::from_bytes(
            GuestDecodeField::StatusText,
            text.as_bytes(),
            self.decode_limits,
        )?;
        let view = *self
            .resources
            .get(view)
            .map_err(|source| resource_handle_error(PluginHandleKind::View.into(), &source))?;
        let session = self
            .active_update
            .as_mut()
            .map(WasmtimeActiveUpdateState::session_mut)
            .ok_or(PluginHostImportError::SessionMissing)?;
        session.push_status_text(view, text.as_str())
    }

    /// Observes through table-local resources and returns owned text.
    fn buffer_observe_from_guest(
        &self,
        buffer: &Resource<PluginBufferHandle>,
    ) -> Result<wit_host::BufferSnapshot, PluginHostImportError> {
        let buffer = *self
            .resources
            .get(buffer)
            .map_err(|source| resource_handle_error(PluginHandleKind::Buffer.into(), &source))?;
        let session = self
            .active_update
            .as_ref()
            .map(WasmtimeActiveUpdateState::session)
            .ok_or(PluginHostImportError::SessionMissing)?;
        let snapshot =
            session.observe_buffer(buffer, self.decode_limits.max_payload_byte_limit())?;
        Ok(wit_host::BufferSnapshot {
            revision: snapshot.revision().get(),
            text: snapshot.text().to_owned(),
        })
    }

    /// Decodes guest edits before they can become host proposals.
    fn buffer_propose_edit_from_guest(
        &mut self,
        buffer: &Resource<PluginBufferHandle>,
        edit: wit_host::BufferEdit,
    ) -> Result<(), PluginHostImportError> {
        let edit = decode_buffer_edit(edit, self.decode_limits)?;
        let buffer = *self
            .resources
            .get(buffer)
            .map_err(|source| resource_handle_error(PluginHandleKind::Buffer.into(), &source))?;
        let session = self
            .active_update
            .as_mut()
            .map(WasmtimeActiveUpdateState::session_mut)
            .ok_or(PluginHostImportError::SessionMissing)?;
        session.propose_buffer_edit(buffer, edit.into_buffer_edit())
    }

    /// Validates path and payload before deferred workspace I/O.
    fn workspace_artifact_write_from_guest(
        &mut self,
        path: &str,
        bytes: &[u8],
    ) -> Result<(), PluginHostImportError> {
        let path = DecodedWorkspacePath::from_bytes(
            GuestDecodeField::WorkspaceArtifactWritePath,
            path.as_bytes(),
            self.decode_limits,
        )?;
        let payload = DecodedPayload::from_bytes(
            GuestDecodeField::WorkspaceArtifactWriteBytes,
            bytes,
            self.decode_limits,
        )?;
        let session = self
            .active_update
            .as_mut()
            .map(WasmtimeActiveUpdateState::session_mut)
            .ok_or(PluginHostImportError::SessionMissing)?;
        session.queue_workspace_artifact_write(path.as_ref(), payload.into_bytes())
    }

    /// Starts a host-owned workspace observation task.
    fn workspace_observe_start_from_guest(
        &mut self,
        path: &str,
    ) -> Result<Resource<PluginWorkspaceObserveTaskHandle>, PluginHostImportError> {
        let path = DecodedWorkspacePath::from_bytes(
            GuestDecodeField::WorkspaceObservePath,
            path.as_bytes(),
            self.decode_limits,
        )?;
        let session = self
            .active_update
            .as_mut()
            .map(WasmtimeActiveUpdateState::session_mut)
            .ok_or(PluginHostImportError::SessionMissing)?;
        let resource_authority = session.queue_workspace_observe_task(
            path.as_ref(),
            self.workspace_observe_tasks.queue_mut(),
        )?;
        self.workspace_observe_tasks
            .push_resource_started_in_update(&mut self.resources, resource_authority)
    }

    /// Polls a workspace observation task without exposing retained bytes.
    fn workspace_observe_poll_from_guest(
        &mut self,
        task: &Resource<PluginWorkspaceObserveTaskHandle>,
    ) -> Result<wit_host::WorkspaceObserveTaskPoll, PluginHostImportError> {
        let session = self
            .active_update
            .as_ref()
            .map(WasmtimeActiveUpdateState::session)
            .ok_or(PluginHostImportError::SessionMissing)?;
        let handle = self.workspace_observe_tasks.resolve_authorized_resource(
            &self.resources,
            session,
            task,
        )?;
        if session.has_pending_workspace_observe_task(&handle) {
            return Ok(wit_host::WorkspaceObserveTaskPoll::Pending);
        }
        Ok(workspace_observe_poll_to_guest(
            self.workspace_observe_tasks.poll(&handle),
        ))
    }

    /// Takes a completed workspace observation task result.
    fn workspace_observe_take_from_guest(
        &mut self,
        task: &Resource<PluginWorkspaceObserveTaskHandle>,
    ) -> Result<wit_host::WorkspaceObserveTaskTake, PluginHostImportError> {
        let session = self
            .active_update
            .as_ref()
            .map(WasmtimeActiveUpdateState::session)
            .ok_or(PluginHostImportError::SessionMissing)?;
        let handle = self.workspace_observe_tasks.resolve_authorized_resource(
            &self.resources,
            session,
            task,
        )?;
        if session.has_pending_workspace_observe_task(&handle) {
            return Ok(wit_host::WorkspaceObserveTaskTake::Pending);
        }
        Ok(workspace_observe_take_to_guest(
            self.workspace_observe_tasks.take(&handle),
        ))
    }

    /// Cancels queue-owned task state and stale task resources for a revoked identity.
    pub(super) fn cancel_workspace_observe_tasks(
        &mut self,
        identity: &PluginIdentity,
    ) -> PluginWorkspaceObserveTaskCancellationReport {
        if let Some(active_update) = self.active_update.as_mut() {
            let _removed = active_update
                .session_mut()
                .cancel_pending_workspace_observe_tasks(identity);
        }
        self.workspace_observe_tasks.cancel_identity(identity)
    }

    /// Cancels all runtime-owned task state for host shutdown.
    pub(super) fn cancel_all_workspace_observe_tasks(
        &mut self,
    ) -> Vec<PluginWorkspaceObserveTaskCancellationReport> {
        if let Some(active_update) = self.active_update.as_mut() {
            let identity = active_update.session().identity_proof().clone();
            let _removed = active_update
                .session_mut()
                .cancel_pending_workspace_observe_tasks(&identity);
        }
        self.workspace_observe_tasks.cancel_all()
    }

    /// Records an import rejection before returning a redacted trap.
    pub(super) fn reject_import(
        &mut self,
        import: PluginOperationalImport,
        error: PluginHostImportError,
    ) -> wasmtime::Error {
        if self.import_rejection.is_none() {
            self.import_rejection = Some(WasmtimeHostImportRejection {
                import,
                source: error,
            });
        }
        wasmtime::Error::msg("plugin host import rejected")
    }
}

/// First import rejection retained for redacted outer diagnostics.
#[derive(Debug)]
pub(super) struct WasmtimeHostImportRejection {
    /// Safe import label.
    pub(super) import: PluginOperationalImport,
    /// Internal rejection cause.
    pub(super) source: PluginHostImportError,
}

/// Store-local state for one active guest update.
#[derive(Debug)]
struct WasmtimeActiveUpdateState {
    /// Discardable host import session.
    session: PluginHostImportSession,
    /// Rooted resources passed into the guest call.
    resources: WasmtimeUpdateResources,
}

impl WasmtimeActiveUpdateState {
    /// Pairs session authority with the update resources that delimit its lifetime.
    const fn new(session: PluginHostImportSession, resources: WasmtimeUpdateResources) -> Self {
        Self { session, resources }
    }

    /// Active import session.
    const fn session(&self) -> &PluginHostImportSession {
        &self.session
    }

    /// Mutable active import session.
    const fn session_mut(&mut self) -> &mut PluginHostImportSession {
        &mut self.session
    }

    /// Splits the active state when guest control returns.
    fn into_parts(self) -> WasmtimeActiveUpdateStateParts {
        WasmtimeActiveUpdateStateParts {
            session: self.session,
            resources: self.resources,
        }
    }
}

/// Active Wasmtime update split after guest control returns.
struct WasmtimeActiveUpdateStateParts {
    /// Import session to seal or discard.
    session: PluginHostImportSession,
    /// Update resources to delete from the table.
    resources: WasmtimeUpdateResources,
}

impl Host for WasmtimeStoreState {
    fn status_publish(
        &mut self,
        view: Resource<PluginViewHandle>,
        text: String,
    ) -> wasmtime::Result<()> {
        self.status_publish_from_guest(&view, &text)
            .map_err(|error| {
                self.reject_import(WitHostImport::StatusPublish.operational_import(), error)
            })
    }

    fn buffer_propose_edit(
        &mut self,
        buffer: Resource<PluginBufferHandle>,
        edit: wit_host::BufferEdit,
    ) -> wasmtime::Result<()> {
        self.buffer_propose_edit_from_guest(&buffer, edit)
            .map_err(|error| {
                self.reject_import(WitHostImport::BufferProposeEdit.operational_import(), error)
            })
    }

    fn buffer_observe(
        &mut self,
        buffer: Resource<PluginBufferHandle>,
    ) -> wasmtime::Result<wit_host::BufferSnapshot> {
        self.buffer_observe_from_guest(&buffer).map_err(|error| {
            self.reject_import(WitHostImport::BufferObserve.operational_import(), error)
        })
    }

    fn workspace_observe_start(
        &mut self,
        path: String,
    ) -> wasmtime::Result<Resource<PluginWorkspaceObserveTaskHandle>> {
        self.workspace_observe_start_from_guest(&path)
            .map_err(|error| {
                self.reject_import(
                    WitHostImport::WorkspaceObserveStart.operational_import(),
                    error,
                )
            })
    }

    fn workspace_observe_poll(
        &mut self,
        task: Resource<PluginWorkspaceObserveTaskHandle>,
    ) -> wasmtime::Result<wit_host::WorkspaceObserveTaskPoll> {
        self.workspace_observe_poll_from_guest(&task)
            .map_err(|error| {
                self.reject_import(
                    WitHostImport::WorkspaceObservePoll.operational_import(),
                    error,
                )
            })
    }

    fn workspace_observe_take(
        &mut self,
        task: Resource<PluginWorkspaceObserveTaskHandle>,
    ) -> wasmtime::Result<wit_host::WorkspaceObserveTaskTake> {
        self.workspace_observe_take_from_guest(&task)
            .map_err(|error| {
                self.reject_import(
                    WitHostImport::WorkspaceObserveTake.operational_import(),
                    error,
                )
            })
    }

    fn workspace_artifact_write(&mut self, path: String, bytes: Vec<u8>) -> wasmtime::Result<()> {
        self.workspace_artifact_write_from_guest(&path, &bytes)
            .map_err(|error| {
                self.reject_import(
                    WitHostImport::WorkspaceArtifactWrite.operational_import(),
                    error,
                )
            })
    }
}

impl HostWorkspaceObserveTask for WasmtimeStoreState {
    fn drop(&mut self, rep: Resource<PluginWorkspaceObserveTaskHandle>) -> wasmtime::Result<()> {
        let dropped = self
            .workspace_observe_tasks
            .delete_resource(&mut self.resources, rep)
            .map_err(|error| {
                self.reject_import(PluginOperationalImport::WorkspaceObserveTaskDrop, error)
            })?;
        match dropped {
            WasmtimeWorkspaceObserveTaskResourceDrop::Active(handle) => {
                if let Some(active_update) = self.active_update.as_mut() {
                    let _removed = active_update
                        .session_mut()
                        .remove_pending_workspace_observe_task(&handle);
                }
            }
            WasmtimeWorkspaceObserveTaskResourceDrop::Inactive => {}
        }
        Ok(())
    }
}

impl HostBufferHandle for WasmtimeStoreState {
    fn drop(&mut self, rep: Resource<PluginBufferHandle>) -> wasmtime::Result<()> {
        let _dropped = self.resources.delete(rep).map_err(|error| {
            self.reject_import(
                PluginOperationalImport::BufferHandleDrop,
                resource_handle_error(PluginHandleKind::Buffer.into(), &error),
            )
        })?;
        Ok(())
    }
}

impl HostViewHandle for WasmtimeStoreState {
    fn drop(&mut self, rep: Resource<PluginViewHandle>) -> wasmtime::Result<()> {
        let _dropped = self.resources.delete(rep).map_err(|error| {
            self.reject_import(
                PluginOperationalImport::ViewHandleDrop,
                resource_handle_error(PluginHandleKind::View.into(), &error),
            )
        })?;
        Ok(())
    }
}

/// Decodes a guest buffer edit into Alma-owned edit vocabulary.
fn decode_buffer_edit(
    edit: wit_host::BufferEdit,
    limits: GuestDecodeLimits,
) -> Result<DecodedBufferEdit, PluginHostImportError> {
    Ok(match edit {
        wit_host::BufferEdit::SetAll(text) => DecodedBufferEdit::set_all(&text, limits)?,
        wit_host::BufferEdit::Insert(insert) => {
            DecodedBufferEdit::insert(insert.byte_index, &insert.text, limits)?
        }
        wit_host::BufferEdit::Replace(replace) => DecodedBufferEdit::replace(
            replace.range.start,
            replace.range.end,
            &replace.text,
            limits,
        )?,
        wit_host::BufferEdit::Delete(range) => {
            DecodedBufferEdit::delete(range.start, range.end, limits)?
        }
    })
}