diff --git a/lib/Vite.php b/lib/Vite.php
index a5916d7..8abcfa8 100644
--- a/lib/Vite.php
+++ b/lib/Vite.php
@@ -2,106 +2,78 @@
 namespace app\lib;
 
 use Piko;
+use Piko\View;
+use RuntimeException;
 
 class Vite
 {
-// For a real-world example check here:
-// https://github.com/wp-bond/bond/blob/master/src/Tooling/Vite.php
-// https://github.com/wp-bond/boilerplate/tree/master/app/themes/boilerplate
+    private View $view;
+    private array $manifest = [];
+    private string $viteHost = '';
+    private bool $dev = false;
 
-// you might check @vitejs/plugin-legacy if you need to support older browsers
-// https://github.com/vitejs/vite/tree/main/packages/plugin-legacy
-
-    /**
-     * Prints all the html entries needed for Vite
-     */
-
-    public static function vite(string $entry): string
+    public function __construct(View $view)
     {
-        return implode("\n", [
-            static::jsTag($entry),
-            static::jsPreloadImports($entry),
-            static::cssTag($entry)
-        ]);
-    }
+        $this->view = $view;
 
-    // Helpers to print tags
-    private static function jsTag(string $entry): string
-    {
-        if (strpos($entry, '.js') === false) {
-            return '';
+        $manifestFile = Piko::getAlias('@webroot/.vite/manifest.json');
+
+        if (!file_exists($manifestFile) && !$this->dev) {
+            throw new RuntimeException("Manifest file not found ($manifestFile)");
         }
 
-        $url = getenv('VITE_ENV') === 'dev'
-            ? getenv('VITE_HOST') . Piko::getAlias('@vite_web/' . $entry)
-            : static::assetUrl($entry);
-
-        if (!$url) {
-            return '';
+        if (file_exists($manifestFile) && !$this->dev) {
+            $manifest = file_get_contents($manifestFile);
+            $this->manifest = json_decode($manifest, true);
         }
 
-        return '<script type="module" crossorigin src="' . $url . '"></script>';
+        $this->dev = getenv('VITE_ENV') === 'dev';
+        $this->viteHost = getenv('VITE_HOST') ?: '';
     }
 
-    private static function jsPreloadImports(string $entry): string
+    public function loadEntry(string $entry): void
     {
-        if (getenv('VITE_ENV') === 'dev') {
-            return '';
+        if ($this->dev) {
+            $this->loadJsFile($entry);
+
+            return;
         }
 
-        $res = '';
-
-        foreach (static::importsUrls($entry) as $url) {
-            $res .= '<link rel="modulepreload" href="'. $url. '">';
+        if (!isset($this->manifest[$entry])) {
+            throw new RuntimeException("Entry {$entry} not found in Vite manifest");
         }
 
-        return $res;
-    }
-
-    private static function cssTag(string $entry): string
-    {
-        // Not needed on dev, it's inject by Vite
-        if (getenv('VITE_ENV') === 'dev' || strpos($entry, '.css') === false) {
-            return '';
+        if (isset($this->manifest[$entry]['file'])) {
+            $this->loadJsFile($this->manifest[$entry]['file']);
         }
 
-        $url =  static::assetUrl($entry);
-
-        if (!$url) {
-            return '';
-        }
-
-        return '<link rel="stylesheet" type="text/css" href="' . $url . '">';
-    }
-
-    // Helpers to locate files
-    private static function getManifest(): array
-    {
-        $content = file_get_contents(Piko::getAlias('@webroot/manifest.json'));
-
-        return json_decode($content, true);
-    }
-
-    private static function assetUrl(string $entry): string
-    {
-        $manifest = static::getManifest();
-
-        return isset($manifest[$entry])
-            ? Piko::getAlias('@web/' . $manifest[$entry]['file'])
-            : '';
-    }
-
-    private static function importsUrls(string $entry): array
-    {
-        $urls = [];
-        $manifest = static::getManifest();
-
-        if (!empty($manifest[$entry]['imports'])) {
-            foreach ($manifest[$entry]['imports'] as $imports) {
-                $urls[] = Piko::getAlias('@web/' . $manifest[$imports]['file']);
+        // Load JS imports
+        foreach ((array)$this->manifest[$entry]['imports'] as $import) {
+            if (isset($this->manifest[$import]['file'])) {
+                $this->loadJsFile($this->manifest[$import]['file'], true);
             }
         }
 
-        return $urls;
+        // Load Stylesheets
+        foreach ((array)$this->manifest[$entry]['css'] as $file) {
+            $this->view->registerCSSFile(Piko::getAlias('@web/' . $file));
+        }
+    }
+
+    private function loadJsFile($file, $import = false): void
+    {
+        if ($import && $this->dev) return;
+
+        if (strpos($file, '.js') === false) return;
+
+        $url = $this->dev ? rtrim($this->viteHost,'/') . '/' . $file : Piko::getAlias('@web/' . $file);
+
+        if ($import) {
+            $this->view->head[] = '<link rel="modulepreload" href="'. $url. '">';
+
+            return;
+        }
+
+        $this->view->endBody[] = '<script type="module" crossorigin src="' . $url . '"></script>';
     }
 }
diff --git a/modules/site/Module.php b/modules/site/Module.php
index a63d186..f5b4928 100644
--- a/modules/site/Module.php
+++ b/modules/site/Module.php
@@ -1,16 +1,14 @@
 <?php
 namespace app\modules\site;
 
-use Piko;
 use Piko\Module\Event\CreateControllerEvent;
 use Piko\Controller\Event\BeforeActionEvent;
+use app\lib\Vite;
 
 class Module extends \Piko\Module
 {
     public function bootstrap()
     {
-        Piko::setAlias('@vite_web', '/dev');
-
         // Instanciate once i18n to setup the language config
         $this->application->getComponent('Piko\I18n');
 
@@ -21,7 +19,9 @@ class Module extends \Piko\Module
         $view = $this->application->getComponent('Piko\View');
         $view->params['user'] = $user;
         $view->params['language'] = $this->application->language;
-        $view->attachBehavior('vite', 'app\lib\Vite::vite');
+        // $view->attachBehavior('vite', 'app\lib\Vite::vite');
+        $vite = new Vite($view);
+        $vite->loadEntry('main.js');
 
         $userModule = $this->application->getModule('user');
         assert ($userModule instanceof \app\modules\user\Module);
diff --git a/modules/site/layouts/main.php b/modules/site/layouts/main.php
index 79bb83a..69c7601 100644
--- a/modules/site/layouts/main.php
+++ b/modules/site/layouts/main.php
@@ -14,7 +14,6 @@ if (!$this->title) $this->title = 'Openai';
     <meta name="viewport" content="width=device-width, initial-scale=1">
     <title><?= $this->escape($this->title) ?></title>
     <?= $this->head() ?>
-    <?= $this->vite('main.css') ?>
 </head>
 <body>
   <nav class="offcanvas offcanvas-start" tabindex="-1" id="offcanvasNavbar" aria-labelledby="offcanvasNavbarLabel">
@@ -63,7 +62,7 @@ if (!$this->title) $this->title = 'Openai';
 			<span class="hamb-bottom"></span>
   </button>
 
-  
+
   <?php if (isset($this->params['message']) && is_array($this->params['message'])): ?>
   <div class="container alert alert-<?= $this->params['message']['type'] ?> alert-dismissible fade show" role="alert">
     <button type="button" class="btn-close" data-bs-dismiss="alert" aria-label="Close"></button>
@@ -73,6 +72,5 @@ if (!$this->title) $this->title = 'Openai';
   <?= $content ?>
 
   <?= $this->endBody() ?>
-  <?= $this->vite('main.js') ?>
 </body>
 </html>
diff --git a/modules/site/layouts/minimal.php b/modules/site/layouts/minimal.php
index a629d60..3d7f107 100644
--- a/modules/site/layouts/minimal.php
+++ b/modules/site/layouts/minimal.php
@@ -11,13 +11,11 @@ if (!$this->title) $this->title = 'Openai';
     <meta name="viewport" content="width=device-width, initial-scale=1">
     <title><?= $this->escape($this->title) ?></title>
     <?= $this->head() ?>
-    <?= $this->vite('main.css') ?>
 </head>
 <body class="d-flex align-items-center py-4">
 
 <?= $content ?>
 
   <?= $this->endBody() ?>
-  <?= $this->vite('main.js') ?>
 </body>
 </html>