2 files changed, 84 insertions, 5 deletions

diff --git a/js/document/html-document.js b/js/document/html-document.js
index ba5eea79..564daef3 100755
--- a/js/document/html-document.js
+++ b/js/document/html-document.js
@@ -388,6 +388,11 @@ exports.HTMLDocument = Montage.create(TextDocument, {
             this._document = this.iframe.contentWindow.document;
             this._window = this.iframe.contentWindow;
             //
+            for (var k in this._document.styleSheets) {
+                this._document.styleSheets[k].ninjatemplate = true;
+                //TODO: Add as attribute
+            }
+            //
             if(!this.documentRoot.Ninja) this.documentRoot.Ninja = {};
             //Inserting user's document into template
             this._templateDocument.head.innerHTML = this._userDocument.content.head;
@@ -423,6 +428,7 @@ exports.HTMLDocument = Montage.create(TextDocument, {
                                                                 cssData = this.application.ninja.coreIoApi.readFile({uri: fileUri});
                                                                 //Creating tag with file content
                                                                 tag = this.iframe.contentWindow.document.createElement('style');
+                                                                tag.setAttribute('type', 'text/css');
                                                                 tag.setAttribute('ninjauri', fileUri);
                                                                 tag.setAttribute('ninjafileurl', cssUrl);
                                                                 tag.setAttribute('ninjafilename', cssUrl.split('/')[cssUrl.split('/').length-1]);
@@ -440,9 +446,11 @@ exports.HTMLDocument = Montage.create(TextDocument, {
                                                         }
                         }
                                         }
+                                        ////////////////////////////////////////////////////////////////////////////
+                                        ////////////////////////////////////////////////////////////////////////////
                                         
                                         //TODO: Revisit this logic
-                                        this._styles = this._document.styleSheets[1];
+                                        /* this._styles = this._document.styleSheets[1]; */
                                         this._stylesheets = this._document.styleSheets; // Entire stlyesheets array
                                         
                                         console.log(this._document.styleSheets);
@@ -564,7 +572,34 @@ exports.HTMLDocument = Montage.create(TextDocument, {
         value: function () {
                 //TODO: Add code view logic and also styles for HTML
                 if (this.currentView === 'design') {
-                        return {mode: 'html', document: this._userDocument, webgl: this.glData, style: this._styles, head: this._templateDocument.head.innerHTML, body: this._templateDocument.body.innerHTML};
+                        var styles = [];
+                        for (var k in this._document.styleSheets) {
+                        if (!this._document.styleSheets[k].ninjatemplate) {
+                                styles.push(this._document.styleSheets[k]);
+                        }
+                }
+                        return {mode: 'html', document: this._userDocument, webgl: this.glData, styles: styles, head: this._templateDocument.head.innerHTML, body: this._templateDocument.body.innerHTML};
+                } else if (this.currentView === "code"){
+                        //TODO: Would this get call when we are in code of HTML?
+                } else {
+                        //Error
+                }
+        }
+        },
+        ////////////////////////////////////////////////////////////////////
+        //
+        saveAll: {
+                enumerable: false,
+        value: function () {
+                //TODO: Add code view logic and also styles for HTML
+                if (this.currentView === 'design') {
+                        var css = [];
+                        for (var k in this._document.styleSheets) {
+                        if (!this._document.styleSheets[k].ninjatemplate) {
+                                css.push(this._document.styleSheets[k]);
+                        }
+                }
+                        return {mode: 'html', document: this._userDocument, webgl: this.glData, css: css, head: this._templateDocument.head.innerHTML, body: this._templateDocument.body.innerHTML};
                 } else if (this.currentView === "code"){
                         //TODO: Would this get call when we are in code of HTML?
                 } else {
diff --git a/js/mediators/io-mediator.js b/js/mediators/io-mediator.js
index 56869839..de50f387 100644
--- a/js/mediators/io-mediator.js
+++ b/js/mediators/io-mediator.js
@@ -202,10 +202,54 @@ exports.IoMediator = Montage.create(Component, {
                 //
                 template.document.content.document.body.innerHTML = template.body;
                 template.document.content.document.head.innerHTML = template.head;
-                //TODO: Remove temp fix for styles
-                if (template.style) {
-                        template.document.content.document.head.getElementsByTagName('style')[0].innerHTML = this.getCssFromRules(template.style.cssRules);
+                //
+                var styletags = template.document.content.document.getElementsByTagName('style'),
+                        linktags = template.document.content.document.getElementsByTagName('link');
+                //
+                for (var j in styletags) {
+                        if (styletags[j].getAttribute) {
+                                if(styletags[j].getAttribute('ninjauri') !== null) {
+                                        try {
+                                                template.document.content.document.head.removeChild(styletags[j]);
+                                        } catch (e) {
+                                                try {
+                                                        template.document.content.document.body.removeChild(styletags[j]);
+                                                } catch (e) {
+                                                        //
+                                                }
+                                        }
+                                        
+                                }
+                        }
+                }
+                //
+                for (var l in linktags) {
+                        if (linktags[l].getAttribute && linktags[l].getAttribute('disabled')) {
+                                linktags[l].removeAttribute('disabled');
+                        }
+                }
+                //
+                if (template.styles) {
+                        //
+                        var styleCounter = 0,
+                                docStyles = template.document.content.document.getElementsByTagName('style');
+                        for(var i in template.styles) {
+                                if (template.styles[i].ownerNode) {
+                                        if (template.styles[i].ownerNode.getAttribute) {
+                                                if (template.styles[i].ownerNode.getAttribute('ninjauri') === null) {
+                                                        //console.log(docStyles[styleCounter], template.styles[i].cssRules);
+                                                        docStyles[styleCounter].innerHTML = this.getCssFromRules(template.styles[i].cssRules);
+                                                        styleCounter++;
+                                                }
+                                        }
+                                }
+                        }
+                        
+                        //template.document.content.document.head.getElementsByTagName('style')[0].innerHTML = this.getCssFromRules(template.style.cssRules);
+                } else if (template.css) {
+                        console.log('Save all css and style');
                 }
+                //
                 return template.document.content.document.documentElement.outerHTML;
         }
     },