CRP File Merger and Optimizer

Mod.cs

/*
 * Copyright 2022 Felix Schmidt
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using ColossalFramework;
using ColossalFramework.IO;
using ColossalFramework.Packaging;
using ColossalFramework.PlatformServices;
using ColossalFramework.UI;
using ICities;
using UnityEngine;

namespace CrapMerger {
    public class Mod : IUserMod {
        public string Name => "CRP File Merger and Optimizer";
        public string Description => "Use through [Options]";

        private const string InputDirKey = "CrapMerger.inputDir";
        private const string OutputDirKey = "CrapMerger.outputDir";
        private const string OutputFileNameKey = "CrapMerger.outputFileName";
        private const string StripSteamPreviewsKey = "CrapMerger.stripSteamPreviews";
        private const string StripContentManagerPreviewsKey = "CrapMerger.stripContentManagerPreviews";
        private const string ShowReportKey = "CrapMerger.showReport";

        private string _inputDir;
        private string _outputDir;
        private string _outputFileName;
        private bool _stripSteamPreviews;
        private bool _stripContentManagerPreviews;
        private bool _showReport;

        //public Package _output;

        private string OutputFilePath => Path.Combine(_outputDir, _outputFileName);

        public void OnEnabled() {
            LoadSettings();
        }

        public void OnSettingsUI(UIHelperBase group) {
            var inputDir = (UITextField)@group.AddTextfield(
                "Input directory (place the .crp files you want to merge in this folder)",
                _inputDir,
                value => {
                    _inputDir = value;
                    PlayerPrefs.SetString(InputDirKey, _inputDir);
                });
            inputDir.width = 740;
            group.AddButton("Show in explorer", () => Utils.OpenInFileBrowser(_inputDir));
            group.AddSpace(16);
            var outputDir = (UITextField)@group.AddTextfield(
                "Output directory",
                _outputDir,
                value => {
                    _outputDir = value;
                    PlayerPrefs.SetString(OutputDirKey, _outputDir);
                });
            outputDir.width = 740;
            var outputFileName = (UITextField)@group.AddTextfield(
                "Output file name",
                _outputFileName,
                value => {
                    _outputFileName = value;
                    PlayerPrefs.SetString(OutputFileNameKey, _outputFileName);
                });
            outputFileName.width = 740;
            group.AddButton("Show in explorer", () => Utils.OpenInFileBrowser(File.Exists(OutputFilePath) ? OutputFilePath : _outputDir));
            group.AddSpace(16);
            UICheckBox stripSteamPreviews = (UICheckBox) group.AddCheckbox("Remove embedded Workshop PreviewImages (recommended, reduces file size)", _stripSteamPreviews,(value) => {
                _stripSteamPreviews = value;
                PlayerPrefs.SetInt(StripSteamPreviewsKey, _stripSteamPreviews ? 1 : 0);
            });
            UICheckBox stripContentManagerPreviews = (UICheckBox) group.AddCheckbox("Remove Content Manager preview images (reduces file size)", _stripContentManagerPreviews, (value) => {
                _stripContentManagerPreviews = value;
                PlayerPrefs.SetInt(StripContentManagerPreviewsKey, _stripContentManagerPreviews ? 1 : 0);
            });
            UICheckBox showReport = (UICheckBox)group.AddCheckbox("Show report on completion", _showReport, (value) => {
                _showReport = value;
                PlayerPrefs.SetInt(ShowReportKey, _showReport ? 1 : 0);
            });
            group.AddSpace(16);
            group.AddButton("Merge .crp files", Execute);
            group.AddButton("Reset settings", () => {
                DeleteSettings();
                LoadSettings();
                inputDir.text = _inputDir;
                outputDir.text = _outputDir;
                outputFileName.text = _outputFileName;
                stripSteamPreviews.isChecked = _stripSteamPreviews;
                stripContentManagerPreviews.isChecked = _stripContentManagerPreviews;
                showReport.isChecked = _showReport;
            });
        }

        private void DeleteSettings() {
            PlayerPrefs.DeleteKey(InputDirKey);
            PlayerPrefs.DeleteKey(OutputDirKey);
            PlayerPrefs.DeleteKey(OutputFileNameKey);
            PlayerPrefs.DeleteKey(StripSteamPreviewsKey);
            PlayerPrefs.DeleteKey(StripContentManagerPreviewsKey);
            PlayerPrefs.DeleteKey(ShowReportKey);
        }

        private void LoadSettings() {
            _inputDir = PlayerPrefs.GetString(InputDirKey);
            if (string.IsNullOrEmpty(_inputDir)) {
                _inputDir = Path.Combine(DataLocation.localApplicationData, "CrapMergerImport");
                if (!Directory.Exists(_inputDir)) Directory.CreateDirectory(_inputDir);
            }

            _outputDir = PlayerPrefs.GetString(OutputDirKey);
            if (string.IsNullOrEmpty(_outputDir)) _outputDir = DataLocation.assetsPath;

            _outputFileName = PlayerPrefs.GetString(OutputFileNameKey);
            if (string.IsNullOrEmpty(_outputFileName)) _outputFileName = "merged.crp";

            _stripSteamPreviews = PlayerPrefs.GetInt(StripSteamPreviewsKey, 1) == 1;

            _stripContentManagerPreviews = PlayerPrefs.GetInt(StripContentManagerPreviewsKey, 0) == 1;

            _showReport = PlayerPrefs.GetInt(ShowReportKey, 1) == 1;
        }

        private void Execute() {
            try {
                if (!Directory.Exists(_inputDir)) {
                    ShowError("Input directory does not exist!");
                    return;
                }

                if (!Directory.Exists(_outputDir)) {
                    ShowError("Output directory does not exist!");
                    return;
                }

                var files = Directory.GetFiles(_inputDir, "*.crp");
                if (files.Length == 0) {
                    ShowError("No .crp files found in input directory!");
                    return;
                }

                PackageManager.DisableEvents();
                var outputPackage = new Package(Path.GetFileNameWithoutExtension(_outputFileName));
                if (PlatformService.active) outputPackage.packageAuthor = "steamid:" + PlatformService.user.userID.AsUInt64;

                var packages = new List<Package>();
                var outdatedPackages = new List<Package>();
                var strippedAssets = new HashSet<Package.Asset>();

                long inputSize = 0;
                foreach (var path in files) {
                    inputSize += new FileInfo(path).Length;
                    var package = new Package(Path.GetFileNameWithoutExtension(path), path, true);
                    if (package.version == outputPackage.version) {
                        packages.Add(package);
                        if (_stripSteamPreviews || _stripContentManagerPreviews) {
                            foreach (var assetInfo in package.FilterAssets(UserAssetType.CustomAssetMetaData)) {
                                CustomAssetMetaData customAssetMetaData = assetInfo.Instantiate<CustomAssetMetaData>();
                                if(_stripSteamPreviews) strippedAssets.Add(customAssetMetaData.steamPreviewRef);
                                if (_stripContentManagerPreviews) strippedAssets.Add(customAssetMetaData.imageRef);
                            }
                        }
                    } else {
                        outdatedPackages.Add(package);
                    }
                }
                
                if (outdatedPackages.Count > 0) {
                    var builder = new StringBuilder("Cannot merge files that were created with a previous version of them game! Please load and save these assets with the asset editor and try again!\n\n");
                    foreach (var package in outdatedPackages) builder.AppendLine(Path.GetFileName(package.packagePath));
                    ShowError(builder.ToString());
                    return;
                }

                List<ReportRow> reportRows = new List<ReportRow>();
                foreach (var package in packages) {
                    Debug.Log($"Package {package.packageName} (main asset: {package.packageMainAsset})");
                    foreach (var assetObj in package) {
                        var assetInfo = (Package.Asset)assetObj;
                        Debug.Log($"├─── Asset {assetInfo.name} ({assetInfo.type}), checksum: {assetInfo.checksum}");

                        if (strippedAssets.Contains(assetInfo)) {
                            Debug.Log($"Stripping asset {assetInfo.name}...");
                            reportRows.Add(new ReportRow {
                                asset = assetInfo,
                                stripped = true
                            });
                            continue;
                        }

                        if (assetInfo.type == Package.AssetType.Texture ||
                            assetInfo.type == Package.AssetType.StaticMesh) {

                            var existingInfo = outputPackage.FindByChecksum(assetInfo.checksum);
                            if (existingInfo != null) {
                                if (!SequenceEquals(GetAssetData(assetInfo), existingInfo.data)) {
                                    ShowError("Checksum collision! Found 2 assets with the same checksum that contain different data!");
                                    return;
                                }

                                Debug.Log($"Found duplicate asset in output package. Stripping {assetInfo.name}...");
                                reportRows.Add(new ReportRow {
                                    asset = assetInfo,
                                    replacement = existingInfo,
                                    stripped = true
                                });
                                continue; // skip this asset!
                            }
                        } else if (assetInfo.type == UserAssetType.CustomAssetMetaData) {
                            if (_stripSteamPreviews || _stripContentManagerPreviews) {
                                var metaData = assetInfo.Instantiate<CustomAssetMetaData>();
                                if (_stripSteamPreviews) {
                                    metaData.steamPreviewRef = null;
                                    Debug.Log("Removed steamPreviewRef");
                                }
                                if (_stripContentManagerPreviews) {
                                    metaData.imageRef = null;
                                    Debug.Log("Removed imageRef");
                                }
                                var outputInfo = outputPackage.AddAsset(assetInfo.name, metaData, UserAssetType.CustomAssetMetaData);
                                reportRows.Add(new ReportRow {
                                    asset = assetInfo,
                                    replacement = outputInfo,
                                    modified = true
                                });
                               
                                continue;
                            }
                        }

                        var outputInfo2 = outputPackage.AddAsset(new Package.Asset(assetInfo.name, GetAssetData(assetInfo), assetInfo.type, false));
                        reportRows.Add(new ReportRow {
                            asset = assetInfo,
                            replacement = outputInfo2,
                        });
                        
                    }
                }

                PackageManager.EnabledEvents();
                outputPackage.Save(OutputFilePath);

                if (_showReport) {
                    long outputSize = new FileInfo(OutputFilePath).Length;
                    ShowReport(reportRows, inputSize, outputSize);
                }

                //_output = new Package(null, OutputFilePath, true); // for debugging

            } catch (Exception e) {
                ShowError(e.ToString());
            }
        }

        public byte[] GetAssetData(Package.Asset assetInfo) {
            using (Stream stream = assetInfo.GetStream()) {
                using (var ms = new MemoryStream()) {
                    stream.CopyTo(ms, assetInfo.size);
                    return ms.ToArray();
                }
            }
        }

        public bool SequenceEquals(byte[] a, byte[] b) {
            if (a.Length != b.Length) return false;
            for (var i = 0; i < a.Length; i++) {
                if (a[i] != b[i]) return false;
            }
            return true;
        }

        private void ShowError(string message) {
            UIView.library.ShowModal<ExceptionPanel>("ExceptionPanel").SetMessage(Name, message,true);
        }

        private void ShowReport(List<ReportRow> rows, long inputSize, long outputSize) {
            var builder = new StringBuilder(@"
<!doctype html>
<html>
<head>
  <meta charset='utf-8'>
  <title>Crap Merger Report</title>
  <style>
    body {
      background-color: white;
      font-family: sans-serif;
    }
    table {
      width: 100%;
      border-collapse: collapse;
      border: 1px solid black;
    }
    th, td {
      text-align: left;
      padding: 4px;
      border: 1px solid black;
    }
    .package {
      padding-top: 20px;
      font-weight: bold;
    }
    .stripped {
      font-weight: bold;
      color: crimson;
    }
    .duplicate {
      font-weight: bold;
      color: orange;
    }
    .copied {
      font-weight: bold;
      color: green;
    }
  </style>
</head>
<body>
  <h1>Crap Merger Report</h1>"
            );

            builder.Append($@"
  <p>Input File Size: {BytesToString(inputSize)}</p>
  <p>Output File Size: {BytesToString(outputSize)}</p>"
            );
            builder.Append(@"
  <table>
    <tr>
      <th>Asset</th><th>Type</th><th>Size</th><th>Status</th>
    </tr>"
            );


            string lastPackage = null;
            foreach (var row in rows) {
                var package = Path.GetFileName(row.asset.package.packagePath);

                if (lastPackage != package) {
                    lastPackage = package;
                    builder.Append($@"
    <tr>
      <td colspan='4' class='package'>{package}</td>
    </tr>"
                    );
                }

                string status;
                if (row.stripped) {
                    if (row.replacement != null) {
                        status = $"<td class='duplicate'>Stripped duplicate</td>";
                    } else {
                        status = "<td class='stripped'>Stripped</td>";
                    }
                } else {
                    if (row.modified) {
                        status = "<td class='copied'>Copied with changes</td>";
                    } else {
                        status = "<td class='copied'>Copied</td>";
                    }
                }

                int size;
                if (row.replacement != null) {
                    size = row.replacement.data.Length;
                } else {
                    size = row.asset.size;
                }

                builder.Append($@"
    <tr>
      <td>{row.asset.name}</td><td>{row.asset.type}</td><td>{BytesToString(size)}</td>{status}
    </tr>"
                );
            }

            builder.Append(@"
  </table>
</body>
</html>"
            );

            var reportPath = Path.Combine(_inputDir, "CrapMergerReport.html");
            File.WriteAllText(reportPath, builder.ToString());
            FileSystemUtils.OpenURLInOverlayOrBrowser(reportPath);
        }

        static string BytesToString(long byteCount) {
            string[] suf = { " B", " KB", " MB", " GB", " TB", " PB", " EB" };
            if (byteCount == 0)
                return "0" + suf[0];
            long bytes = Math.Abs(byteCount);
            int place = Convert.ToInt32(Math.Floor(Math.Log(bytes, 1024)));
            double num = Math.Round(bytes / Math.Pow(1024, place), 1);
            return (Math.Sign(byteCount) * num) + suf[place];
        }

        private class ReportRow {
            public Package.Asset asset;
            public Package.Asset replacement;
            public bool stripped;
            public bool modified;
        }
    }
}