1 files changed, 39 insertions, 3 deletions
diff --git a/lib/generators/site.js b/lib/generators/site.js
index 7a2a3aa..6caf372 100644
--- a/lib/generators/site.js
+++ b/lib/generators/site.js
@@ -12,14 +12,50 @@ var pageUtil = require("../utils/page");
 var Generator = function() {
     BaseGenerator.apply(this, arguments);
 
-    this.env = new nunjucks.Environment(
-		new nunjucks.FileSystemLoader(book.root),
+    // Style to integrates i nthe output
+    this.styles = ["website"];
+
+    // base folder for templates
+    this.templatesRoot = path.resolve(__dirname, "../../theme/templates/website")
+};
+util.inherits(Generator, BaseGenerator);
+
+// Prepare the genertor
+BaseGenerator.prototype.load = function() {
+    var that = this;
+
+    return BaseGenerator.prototype.load.apply(this)
+    .then(function() {
+        return that.loadTemplates();
+    });
+};
+
+// Load all styles
+Generator.prototype.loadStyles = function() {
+    var that = this;
+    this.styles = _.chain(this.styles)
+        .map(function(style) {
+            var stylePath = that.options.styles[style];
+            if (fs.existsSync(path.resolve(that.book.root, stylePath))) {
+                return stylePath;
+            }
+            return null;
+        })
+        .compact()
+        .value();
+};
+
+// Load template engine
+Generator.prototype.loadTemplates = function() {
+	this.env = new nunjucks.Environment(
+		new nunjucks.FileSystemLoader(this.templatesRoot),
 		{
 			autoescape: true
 		}
 	);
+
+	return Q();
 };
-util.inherits(Generator, BaseGenerator);
 
 // Ignore some methods
 Generator.prototype.transferFile = function(input) {