<?php

namespace LoftyIDX\includes\extensions;

use Twig\Environment;
use Twig\TwigFunction;
use LoftyIDX\includes\LoftyIDXHelpers;

defined('ABSPATH') || exit;

class LoftyIDXTwigScriptExtension extends \Twig\Extension\AbstractExtension
{
  private $scripts = [];
  private $handle_counter = 0;
  private $twig;
  private $scripts_removed = false;

  public function __construct(Environment $twig)
  {
    $this->twig = $twig;
    // Don't remove scripts globally in constructor
    // Only remove when actually using vite_js function
  }

  public function getFunctions()
  {
    return [
      new TwigFunction('vite_js', [$this, 'addViteScript'], ['needs_context' => true, 'needs_environment' => true]),
      new TwigFunction('render_scripts', [$this, 'renderScripts']),
    ];
  }

  /**
   * Remove default WordPress script output only when needed
   * This should only affect pages that use vite_js function
   */
  private function removeDefaultScripts()
  {
    if (!$this->scripts_removed) {
      // Remove default WordPress scripts output from wp_head
      // Note: wp_head() calls wp_print_head_scripts(), not wp_print_scripts()
      remove_action('wp_head', 'wp_print_head_scripts');
      $this->scripts_removed = true;
    }
  }

  public function get_vite_manifest()
  {
    static $manifest = null;

    if ($manifest === null) {
      $manifest_path = LOFTY_IDX_PATH . '/assets/.vite/manifest.json';

      global $wp_filesystem;
      if (empty($wp_filesystem)) {
        require_once(ABSPATH . 'wp-admin/includes/file.php');
        WP_Filesystem();
      }

      if (!$wp_filesystem->exists($manifest_path)) {
        throw new \Exception('Vite manifest.json not found. Please run build first.');
      }

      $manifest_content=$wp_filesystem->get_contents($manifest_path);
      $manifest = json_decode($manifest_content, true);

      if ($manifest === null) {
        throw new \Exception('Vite manifest.json is invalid.');
      }
    }

    return $manifest;
  }

  // add type="module" script, for vite esm js
  public function addViteScript(Environment $env, $context, $path, $dependencies = [], $version = '', $args = ['strategy' => 'defer'])
  {
    // Only remove default scripts when we're actually using vite scripts
    $this->removeDefaultScripts();
    
    // Generate a unique handle for the style
    $handle = 'twig-vite-' . ++$this->handle_counter;
    // check if dev mode
    if (LOFTY_IS_DEV) {
      // handle final url based on $path
      $url = esc_url(LoftyIDXHelpers::loftyVites($path));

      // Register the style with WordPress
      wp_register_script($handle, $url, $dependencies, $version, $args);
      wp_enqueue_script($handle);

      // vite js will auto inject related css/scss links

      add_filter('script_loader_tag',  function ($tag, $handle) {
        if (strpos($handle, 'vite') > -1) {  // The type="module" of the introduction of the VITE project need to be added
          return str_replace('<script', '<script type="module"', $tag);
        }
        return $tag;  // The original tag is returned, and other scripts are not modified
      }, 10, 2);

      // Store the handle for our custom output
      $this->scripts[] = $handle;
    } else {
      // production mode
      $manifest = $this->get_vite_manifest();
      // Get the manifest entry
      if (!isset($manifest[$path])) {
        throw new \Exception(esc_html("Entry {$path} not found in Vite manifest."));
      }

      $this->handleManifestPath($manifest, $path, $dependencies, $version, $args);
    }

    return '';
  }

  public function handleManifestPath($manifest=[], $path='', $dependencies = [], $version = '', $args = ['strategy' => 'defer'])
  {
    $manifest_data=$manifest[$path];

    $handle_prefix = 'twig-vite-';
    $handle=$handle_prefix . $manifest_data['file'];

    // import all chunk js files related chunk css files
    if(isset($manifest_data['imports']) && is_array($manifest_data['imports'])){
      // get chunk css files by imports js files
      foreach($manifest_data['imports'] as $chunk_js){
        if(isset($manifest[$chunk_js]) && isset($manifest[$chunk_js]['css']) 
        && is_array($manifest[$chunk_js]['css'])){
          $chunk_css=$manifest[$chunk_js]['css'];
          foreach($chunk_css as $css_file){
            wp_enqueue_style(
              $handle_prefix . $css_file,
              LOFTY_IDX_URL . 'assets/' . $css_file,
              [],
              $version
            );
          }
        }
      }
    }

    // Add CSS files from manifest
    if (isset($manifest_data['css']) && is_array($manifest_data['css'])) {
      foreach ($manifest_data['css'] as $css_file) {
        wp_enqueue_style(
          $handle_prefix . $css_file,
          LOFTY_IDX_URL . 'assets/' . $css_file,
          [],
          $version
        );
      }
    }

    // path entry js
    wp_register_script(
      $handle, 
      LOFTY_IDX_URL . 'assets/' . $manifest_data['file'],
      $dependencies, 
      $version, 
      $args
    );
    wp_enqueue_script($handle);

    add_filter('script_loader_tag',  function ($tag, $handle) {
      if (strpos($handle, 'vite') > -1) {
        return str_replace('<script', '<script type="module"', $tag);
      }
      return $tag;
    }, 10, 2);

    // handle dynamic imports
    if(isset($manifest_data['dynamicImports']) && is_array($manifest_data['dynamicImports'])){
      foreach($manifest_data['dynamicImports'] as $dynamic_path){
        $this->handleManifestPath($manifest, $dynamic_path, $dependencies, $version, $args);
      }
    }
  }

  /**
   * @all whether output all styles
   */
  public function renderScripts()
  {
    global $wp_scripts;

    if (empty($this->scripts)) {
      return;
    }

    $handles = array_filter($this->scripts, function ($handle) use ($wp_scripts) {
      return isset($wp_scripts->registered[$handle]);
    });

    wp_print_scripts($handles);
    // Clear the scripts after rendering
    $this->scripts = [];
  }
}