dprint-core 0.71.1

use std::borrow::Cow;
use std::cell::OnceCell;
use std::cell::RefCell;
use std::path::PathBuf;

use crate::configuration::ConfigKeyMap;
use crate::plugins::CancellationToken;
use crate::plugins::CheckConfigUpdatesMessage;
use crate::plugins::FormatConfigId;
use crate::plugins::FormatError;
use crate::plugins::FormatRange;
use crate::plugins::FormatResult;
use crate::plugins::PluginResolveConfigurationResult;
use crate::plugins::RawFormatConfig;
use crate::plugins::SyncFormatRequest;
use crate::plugins::SyncHostFormatRequest;
use crate::plugins::SyncPluginHandler;

use super::JsonResponse;

/// State of a Wasm plugin and the implementation of the functions
/// it exports, which are generated by `generate_plugin_code!`.
///
/// This lives in dprint-core rather than in the macro so that it's
/// compiled and linted here instead of in each plugin's crate.
#[doc(hidden)]
pub struct WasmPluginState<TPlugin, TConfiguration: Clone + serde::Serialize> {
  create_plugin: fn() -> TPlugin,
  plugin: OnceCell<RefCell<TPlugin>>,
  shared_bytes: RefCell<Vec<u8>>,
  override_config: RefCell<Option<ConfigKeyMap>>,
  file_path: RefCell<Option<PathBuf>>,
  formatted_text: RefCell<Option<Vec<u8>>>,
  error_text: RefCell<Option<String>>,
  unresolved_configs: RefCell<ConfigIdMap<RawFormatConfig>>,
  resolved_configs: RefCell<ConfigIdMap<PluginResolveConfigurationResult<TConfiguration>>>,
}

// SAFETY: Wasm plugins are only ever executed on a single thread.
// https://github.com/rust-lang/rust/issues/53639#issuecomment-790091647
unsafe impl<TPlugin, TConfiguration: Clone + serde::Serialize> Sync for WasmPluginState<TPlugin, TConfiguration> {}

impl<TPlugin, TConfiguration: Clone + serde::Serialize> WasmPluginState<TPlugin, TConfiguration> {
  pub const fn new(create_plugin: fn() -> TPlugin) -> Self {
    Self {
      create_plugin,
      plugin: OnceCell::new(),
      shared_bytes: RefCell::new(Vec::new()),
      override_config: RefCell::new(None),
      file_path: RefCell::new(None),
      formatted_text: RefCell::new(None),
      error_text: RefCell::new(None),
      unresolved_configs: RefCell::new(ConfigIdMap::new()),
      resolved_configs: RefCell::new(ConfigIdMap::new()),
    }
  }
}

impl<TPlugin: SyncPluginHandler<TConfiguration>, TConfiguration: Clone + serde::Serialize> WasmPluginState<TPlugin, TConfiguration> {
  // FORMATTING

  pub fn set_override_config(&self) {
    let bytes = self.take_shared_bytes();
    let config = serde_json::from_slice(&bytes).unwrap();
    *self.override_config.borrow_mut() = Some(config);
  }

  pub fn set_file_path(&self) {
    // convert windows back slashes to forward slashes so it works with PathBuf
    let text = self.take_shared_bytes_string().replace('\\', "/");
    *self.file_path.borrow_mut() = Some(PathBuf::from(text));
  }

  pub fn format(&self, config_id: u32, range: FormatRange) -> u8 {
    let config_id = FormatConfigId::from_raw(config_id);
    self.ensure_initialized(config_id);
    let override_config = self.override_config.borrow_mut().take();
    let resolved_configs = self.resolved_configs.borrow();
    let config = match override_config {
      Some(override_config) => Cow::Owned(self.create_resolved_config_result(config_id, override_config).config),
      None => Cow::Borrowed(&resolved_configs.get(config_id).unwrap().config),
    };
    let file_path = self.file_path.borrow_mut().take().expect("Expected the file path to be set.");
    let file_bytes = self.take_shared_bytes();

    let request = SyncFormatRequest::<TConfiguration> {
      file_path: &file_path,
      file_bytes,
      config: &config,
      config_id,
      range,
      token: &HostCancellationToken,
    };
    let result = self.plugin().borrow_mut().format(request, |request| self.format_with_host(request));
    match result {
      Ok(None) => {
        0 // no change
      }
      Ok(Some(formatted_text)) => {
        *self.formatted_text.borrow_mut() = Some(formatted_text);
        1 // change
      }
      Err(err) => {
        *self.error_text.borrow_mut() = Some(err.to_string());
        2 // error
      }
    }
  }

  pub fn get_formatted_text(&self) -> usize {
    let formatted_text = self.formatted_text.borrow_mut().take().expect("Expected to have formatted text.");
    self.set_shared_bytes(formatted_text)
  }

  pub fn get_error_text(&self) -> usize {
    let error_text = self.error_text.borrow_mut().take().expect("Expected to have error text.");
    self.set_shared_bytes(error_text.into_bytes())
  }

  // INFORMATION & CONFIGURATION

  pub fn get_plugin_info(&self) -> usize {
    let plugin_info = self.plugin().borrow_mut().plugin_info();
    self.set_shared_bytes(serde_json::to_vec(&plugin_info).unwrap())
  }

  pub fn get_license_text(&self) -> usize {
    let license_text = self.plugin().borrow_mut().license_text();
    self.set_shared_bytes(license_text.into_bytes())
  }

  pub fn get_resolved_config(&self, config_id: u32) -> usize {
    let bytes = self.with_resolved_config_result(config_id, |result| serde_json::to_vec(&result.config).unwrap());
    self.set_shared_bytes(bytes)
  }

  pub fn get_config_diagnostics(&self, config_id: u32) -> usize {
    let bytes = self.with_resolved_config_result(config_id, |result| serde_json::to_vec(&result.diagnostics).unwrap());
    self.set_shared_bytes(bytes)
  }

  pub fn get_config_file_matching(&self, config_id: u32) -> usize {
    let bytes = self.with_resolved_config_result(config_id, |result| serde_json::to_vec(&result.file_matching).unwrap());
    self.set_shared_bytes(bytes)
  }

  // INITIALIZATION

  pub fn register_config(&self, config_id: u32) {
    let config_id = FormatConfigId::from_raw(config_id);
    let bytes = self.take_shared_bytes();
    let config: RawFormatConfig = serde_json::from_slice(&bytes).unwrap();
    self.unresolved_configs.borrow_mut().insert(config_id, config);
    self.resolved_configs.borrow_mut().remove(config_id); // clear
  }

  pub fn release_config(&self, config_id: u32) {
    let config_id = FormatConfigId::from_raw(config_id);
    self.unresolved_configs.borrow_mut().remove(config_id);
    self.resolved_configs.borrow_mut().remove(config_id);
  }

  pub fn check_config_updates(&self) -> usize {
    let bytes = self.take_shared_bytes();
    let response = match self.try_check_config_updates(&bytes) {
      Ok(value) => JsonResponse::Ok(value),
      Err(err) => JsonResponse::Err(err.to_string()),
    };
    self.set_shared_bytes(serde_json::to_vec(&response).unwrap())
  }

  // LOW LEVEL SENDING AND RECEIVING

  pub fn get_shared_bytes_ptr(&self) -> *const u8 {
    self.shared_bytes.borrow().as_ptr()
  }

  pub fn clear_shared_bytes(&self, size: usize) -> *const u8 {
    let mut shared_bytes = self.shared_bytes.borrow_mut();
    *shared_bytes = vec![0; size];
    shared_bytes.as_mut_ptr()
  }

  fn plugin(&self) -> &RefCell<TPlugin> {
    self.plugin.get_or_init(|| RefCell::new((self.create_plugin)()))
  }

  fn format_with_host(&self, request: SyncHostFormatRequest) -> FormatResult {
    let file_path = request.file_path.to_string_lossy();
    let override_config = if !request.override_config.is_empty() {
      Cow::Owned(serde_json::to_string(request.override_config).unwrap())
    } else {
      Cow::Borrowed("")
    };
    let range = request.range.unwrap_or(0..request.file_bytes.len());

    // SAFETY: the pointers and lengths are valid for the duration of the call
    let result = unsafe {
      host::host_format(
        file_path.as_ptr(),
        file_path.len() as u32,
        range.start as u32,
        range.end as u32,
        override_config.as_ptr(),
        override_config.len() as u32,
        request.file_bytes.as_ptr(),
        request.file_bytes.len() as u32,
      )
    };
    match result {
      0 => {
        // no change
        Ok(None)
      }
      1 => {
        // change
        let length = unsafe { host::host_get_formatted_text() };
        Ok(Some(self.get_bytes_from_host(length)))
      }
      2 => {
        // error
        let length = unsafe { host::host_get_error_text() };
        let error_text = String::from_utf8(self.get_bytes_from_host(length)).unwrap();
        Err(FormatError::new(error_text))
      }
      value => panic!("unknown host format value: {}", value),
    }
  }

  fn get_bytes_from_host(&self, length: u32) -> Vec<u8> {
    let ptr = self.clear_shared_bytes(length as usize);
    // SAFETY: the host writes `length` bytes to the buffer that was just allocated
    unsafe { host::host_write_buffer(ptr) };
    self.take_shared_bytes()
  }

  fn with_resolved_config_result<R>(&self, config_id: u32, func: impl FnOnce(&PluginResolveConfigurationResult<TConfiguration>) -> R) -> R {
    let config_id = FormatConfigId::from_raw(config_id);
    self.ensure_initialized(config_id);
    func(self.resolved_configs.borrow().get(config_id).unwrap())
  }

  fn ensure_initialized(&self, config_id: FormatConfigId) {
    if !self.resolved_configs.borrow().contains(config_id) {
      let config_result = self.create_resolved_config_result(config_id, ConfigKeyMap::new());
      self.resolved_configs.borrow_mut().insert(config_id, config_result);
    }
  }

  fn create_resolved_config_result(&self, config_id: FormatConfigId, override_config: ConfigKeyMap) -> PluginResolveConfigurationResult<TConfiguration> {
    let unresolved_configs = self.unresolved_configs.borrow();
    let Some(config) = unresolved_configs.get(config_id) else {
      panic!("Plugin must have config set before use (id: {:?}).", config_id);
    };
    let mut plugin_config = config.plugin.clone();
    for (key, value) in override_config {
      plugin_config.insert(key, value);
    }
    self.plugin().borrow_mut().resolve_config(plugin_config, &config.global)
  }

  fn try_check_config_updates(&self, bytes: &[u8]) -> Result<serde_json::Value, FormatError> {
    let message: CheckConfigUpdatesMessage = serde_json::from_slice(bytes)?;
    let result = self.plugin().borrow().check_config_updates(message)?;
    Ok(serde_json::to_value(&result)?)
  }

  fn take_shared_bytes_string(&self) -> String {
    String::from_utf8(self.take_shared_bytes()).unwrap()
  }

  fn take_shared_bytes(&self) -> Vec<u8> {
    std::mem::take(&mut *self.shared_bytes.borrow_mut())
  }

  fn set_shared_bytes(&self, bytes: Vec<u8>) -> usize {
    let length = bytes.len();
    *self.shared_bytes.borrow_mut() = bytes;
    length
  }
}

#[derive(Debug)]
struct HostCancellationToken;

impl CancellationToken for HostCancellationToken {
  fn is_cancelled(&self) -> bool {
    unsafe { host::host_has_cancelled() == 1 }
  }
}

/// A map keyed by config id.
///
/// A plugin only holds a few configs at a time, so a vec is used
/// instead of a hash map in order to keep the plugin small.
struct ConfigIdMap<T>(Vec<(FormatConfigId, T)>);

impl<T> ConfigIdMap<T> {
  const fn new() -> Self {
    Self(Vec::new())
  }

  fn get(&self, id: FormatConfigId) -> Option<&T> {
    self.0.iter().find(|(key, _)| *key == id).map(|(_, value)| value)
  }

  fn contains(&self, id: FormatConfigId) -> bool {
    self.get(id).is_some()
  }

  fn insert(&mut self, id: FormatConfigId, value: T) {
    match self.0.iter_mut().find(|(key, _)| *key == id) {
      Some((_, existing)) => *existing = value,
      None => self.0.push((id, value)),
    }
  }

  fn remove(&mut self, id: FormatConfigId) {
    self.0.retain(|(key, _)| *key != id);
  }
}

mod host {
  #[link(wasm_import_module = "dprint")]
  unsafe extern "C" {
    pub fn host_has_cancelled() -> i32;
    pub fn host_write_buffer(pointer: *const u8);
    pub fn host_format(
      file_path_ptr: *const u8,
      file_path_len: u32,
      start_range: u32,
      end_range: u32,
      override_config_ptr: *const u8,
      override_config_len: u32,
      file_text_ptr: *const u8,
      file_text_len: u32,
    ) -> u8;
    pub fn host_get_formatted_text() -> u32;
    pub fn host_get_error_text() -> u32;
  }
}