using System;
using System.Collections.Generic;
using System.IO;

using UnityEditor;
using UnityEngine;

namespace Framework.Terrain {

  /// <summary>Imports terrain from various sources</summary>
  [CustomEditor(typeof(TerrainImportSettings))]
  public class TerrainImportEditor : Editor {

    /// <summary>Available color channels in a standard 32 bit texture</summary>
    private readonly static string[] ColorChannels = new string[] {
      "None", "R", "G", "B", "A", "R2", "G2", "B2", "A2"
    };

    /// <summary>Actions the terrain importer/exporter can perform</summary>
    private enum ImportExportAction {
      /// <summary>Do nothing</summary>
      None,
      /// <summary>Import terrain heights from height map</summary>
      ImportHeightMap,
      /// <summary>Export terrain heights to height map</summary>
      ExportHeightMap,
      /// <summary>Import splat maps as terrain layer weights</summary>
      ImportSplatMaps,
      /// <summary>Export up to four terrain layer weights as a splat map</summary>
      ExportPrimarySplatMap,
      /// <summary>Export up to four additional terrain layer weights as a splat map</summary>
      ExportSecondarySplatMap
    };

    /// <summary>Draws the terrain import tool panel in the Inspector</summary>
    public override void OnInspectorGUI() {
      var terrainSettings = base.target as TerrainImportSettings;
      if(terrainSettings == null) {
        GUILayout.Label(
          "Terrain import settings not displayed because " +
          "TerrainImportSettings component not present"
        );
        return;
      }

      UnityEngine.Terrain terrain = terrainSettings.GetComponent<UnityEngine.Terrain>();
      if(terrain == null) {
        GUILayout.Label(
          "Terrain import settings not displayed because Terrain component is missing"
        );
        return;
      }

      // Show the UI for setting up height map import/export
      doHeightMapGUI(terrainSettings);
      ImportExportAction heightMapAction = doHeightMapButtonGUI();
      displayWarningsForMissingHeightTexture(terrainSettings);

      EditorGUILayout.Separator();

      // Show the UI for settings up splat map import/export
      doSplatLayerGUI(terrain, terrainSettings);
      ImportExportAction splatMapAction = doSplatMapButtonGUI();
      displayWarningsForMissingSplatTextures(terrainSettings);

      // If the user clicked one of the buttons, do the appropriate action
      if(heightMapAction == ImportExportAction.ImportHeightMap) {
        importHeightMapFromTexture(terrainSettings, terrain.terrainData);
      } else if(heightMapAction == ImportExportAction.ExportHeightMap) {
        exportHeightMapToTexture(terrain.terrainData);
      } else if(splatMapAction == ImportExportAction.ImportSplatMaps) {
        importSplatMapsFromTexture(terrainSettings, terrain.terrainData);
      } else if(splatMapAction == ImportExportAction.ExportPrimarySplatMap) {
        exportPrimarySplatMapToTexture(terrain.terrainData, terrainSettings.SplatColorIndices);
      } else if(splatMapAction == ImportExportAction.ExportSecondarySplatMap) {
        exportSecondarySplatMapToTexture(terrain.terrainData, terrainSettings.SplatColorIndices);
      }
    }

    /// <summary>Displays the inspector UI for the height map import settings</summary>
    /// <param name="terrainSettings">Terrain settings that will be modified</param>
    private void doHeightMapGUI(TerrainImportSettings terrainSettings) {
      const bool AllowSceneObjects = false;

      using(var horizontalScope = new EditorGUILayout.HorizontalScope()) {
        EditorGUILayout.PrefixLabel("Height Map Input");
        terrainSettings.HeightMapTexture = (Texture2D) EditorGUILayout.ObjectField(
          terrainSettings.HeightMapTexture, typeof(Texture2D), AllowSceneObjects
        );
      }

      using(var horizontalScope = new EditorGUILayout.HorizontalScope()) {
        EditorGUILayout.PrefixLabel("Minimum Height");
        terrainSettings.MinHeight = EditorGUILayout.FloatField(terrainSettings.MinHeight);
      }

      using(var horizontalScope = new EditorGUILayout.HorizontalScope()) {
        EditorGUILayout.PrefixLabel("Maximum Height");
        terrainSettings.MaxHeight = EditorGUILayout.FloatField(terrainSettings.MaxHeight);
      }
    }

    /// <summary>Handles the buttons by which the height map can be imported / exported</summary>
    /// <returns>The action the user wishes to perform, if any</returns>
    private ImportExportAction doHeightMapButtonGUI() {
      var result = ImportExportAction.None;

      using(var horizontalScope = new EditorGUILayout.HorizontalScope()) {
        EditorGUILayout.PrefixLabel(" ");
        using(var verticalScope = new EditorGUILayout.VerticalScope()) {
          if(GUILayout.Button("Import Height Map")) {
            result = ImportExportAction.ImportHeightMap;
          }
          if(GUILayout.Button("Export Height Map")) {
            result = ImportExportAction.ExportHeightMap;
          }
        }
      }

      return result;
    }

    /// <summary>Displays warnings if the height texture is missing</summary>
    /// <param name="terrainSettings">Terrain import settings that will be checked</param>
    private void displayWarningsForMissingHeightTexture(TerrainImportSettings terrainSettings) {
      if(terrainSettings.HeightMapTexture == null) {
        EditorGUILayout.HelpBox("WARNING: No height map texture assigned", MessageType.Warning);
      }
    }

    /// <summary>Displays the UI for setting up splat layer assignments</summary>
    /// <param name="terrain">Terrain from which to take the splat layers</param>
    /// <param name="terrainSettings">Terrain settings that will be modified</param>
    private void doSplatLayerGUI(
      UnityEngine.Terrain terrain, TerrainImportSettings terrainSettings
    ) {
      const bool AllowSceneObjects = false;

      using(var horizontalScope = new EditorGUILayout.HorizontalScope()) {
        EditorGUILayout.PrefixLabel("Splat Map Inputs");
        terrainSettings.SplatMapTexture = (Texture2D) EditorGUILayout.ObjectField(
          terrainSettings.SplatMapTexture, typeof(Texture2D), AllowSceneObjects
        );
        terrainSettings.SecondarySplatMapTexture = (Texture2D) EditorGUILayout.ObjectField(
          terrainSettings.SecondarySplatMapTexture, typeof(Texture2D), AllowSceneObjects
        );
      }

      EditorGUILayout.Separator();

      // If there are no layers, inform the user why he can't set up layer assignments
      int layerCount = terrain.terrainData.alphamapLayers;
      if(layerCount == 0) {
        GUILayout.Label("Add texture layers to your terrain to display layer assignments");
        return;
      }

      // Make sure there are layers set up for each splat layer
      terrainSettings.ResizeSplatColorIndices(layerCount);

      // Let the user choose which color channel to use for each splat layer
      for(int index = 0; index < layerCount; ++index) {
        using(var horizontalScope = new EditorGUILayout.HorizontalScope()) {
          EditorGUILayout.PrefixLabel("Layer " + (index + 1));
          terrainSettings.SplatColorIndices[index] = GUILayout.SelectionGrid(
            terrainSettings.SplatColorIndices[index], ColorChannels, 9
          );
        }
      }

    }

    /// <summary>Draws the splat map buttons in the inspector</summary>
    /// <returns>The action the user wishes to perform, if any</returns>
    private ImportExportAction doSplatMapButtonGUI() {
      var result = ImportExportAction.None;

      using(var horizontalScope = new EditorGUILayout.HorizontalScope()) {
        EditorGUILayout.PrefixLabel(" ");

        using(var verticalScope = new EditorGUILayout.VerticalScope()) {
          if(GUILayout.Button("Import Splat Map(s)")) {
            result = ImportExportAction.ImportSplatMaps;
          }
          using(var innerHorizontalScope = new EditorGUILayout.HorizontalScope()) {
            if(GUILayout.Button("Export Splat Map 1")) {
              result = ImportExportAction.ExportPrimarySplatMap;
            }
            if(GUILayout.Button("Export Splat Map 2")) {
              result = ImportExportAction.ExportSecondarySplatMap;
            }
          }
        }
      }

      return result;
    }

    /// <summary>Displays warnings if required splat textures are missing</summary>
    /// <param name="terrainSettings">Terrain import settings that will be checked</param>
    private void displayWarningsForMissingSplatTextures(TerrainImportSettings terrainSettings) {

      // Generate a string containing all texture warning
      string warnings = string.Empty;
      {
        bool missingPrimarySplatMapTexture =
          (terrainSettings.SplatMapTexture == null) &&
          (terrainSettings.ReferencesPrimarySplatMapTexture);
        if(missingPrimarySplatMapTexture) {
          warnings = "WARNING: no primary splat map texture assigned";
        }
        bool missingSecondarySplatMapTexture =
          (terrainSettings.SecondarySplatMapTexture == null) &&
          (terrainSettings.ReferencesSecondarySplatMapTexture);
        if(missingSecondarySplatMapTexture) {
          if(missingPrimarySplatMapTexture) {
            warnings += "\n";
          }

          warnings += "WARNING: no secondary splat map texture assigned";
        }
      }

      // If any warning messages were triggered, display them
      if(!string.IsNullOrEmpty(warnings)) {
        EditorGUILayout.HelpBox(warnings, MessageType.Warning);
      }

    }

    /// <summary>Returns an array containing all currently selected terrains</summary>
    /// <returns>An array of all selected terrains or null if none were selected</returns>
    private UnityEngine.Terrain[] getSelectedTerrains() {
      GameObject[] selectedGameObjects = UnityEditor.Selection.gameObjects;
      if(selectedGameObjects == null) {
        return null;
      }

      int selectedGameObjectCount = selectedGameObjects.Length;
      var terrains = new List<UnityEngine.Terrain>(capacity: selectedGameObjectCount);

      for(int index = 0; index < selectedGameObjectCount; ++index) {
        UnityEngine.Terrain terrain =
          selectedGameObjects[index].GetComponent<UnityEngine.Terrain>();

        if(terrain != null) {
          terrains.Add(terrain);
        }
      }

      if(terrains.Count == 0) {
        return null;
      } else {
        return terrains.ToArray();
      }
    }

    /// <summary>Imports the terrain's heights from a height map texture</summary>
    /// <param name="importSettings">Import settings containing the texture</param>
    /// <param name="terrainData">Terrain the height map will be imported into</param>
    private void importHeightMapFromTexture(
      TerrainImportSettings importSettings, TerrainData terrainData
    ) {
      if(importSettings.HeightMapTexture == null) {
        EditorUtility.DisplayDialog(
          "Texture not Assigned",
          "There is no height map texture assigned to the import from",
          "OK"
        );
        return;
      }

      TerrainHelper.ImportHeightsFromTexture(terrainData, importSettings.HeightMapTexture);
    }

    /// <summary>Exports the terrain to a height map texture</summary>
    /// <param name="terrainData">Terrain data that will be exported</param>
    private void exportHeightMapToTexture(TerrainData terrainData) {
      Texture2D heightMapTexture = TerrainHelper.GetHeightsAsTexture(terrainData);
      string path = EditorUtility.SaveFilePanel(
        "Select Location to Save Height Map", ".", "heightmap.png", "png"
      );
      if(!string.IsNullOrEmpty(path)) {
        writeToFile(path, heightMapTexture.EncodeToPNG());
      }
    }


    /// <summary>Imports the terrain's splat maps from one or two textures</summary>
    /// <param name="importSettings">Import settings containing the textures</param>
    /// <param name="terrainData">Terrain the splat layer weights will be applied to</param>
    private void importSplatMapsFromTexture(
      TerrainImportSettings importSettings, TerrainData terrainData
    ) {
      var textures = new Texture2D[] {
        importSettings.SplatMapTexture, importSettings.SecondarySplatMapTexture
      };
      TerrainHelper.ImportLayerWeightsFromTextures(
        terrainData, textures, importSettings.SplatColorIndices
      );
    }

    /// <summary>Saves the terrain's primary splat map into a file</summary>
    /// <param name="terrainData">Terrain whose primary splat map will be saved</param>
    /// <param name="splatColorIndices">Color channel assignments for the splat layers</param>
    private void exportPrimarySplatMapToTexture(
      TerrainData terrainData, int[] splatColorIndices
    ) {
      int[] splatLayerIndices = layerIndicesFromColorIndices(splatColorIndices, -1);
      Texture2D splatMapTexture = TerrainHelper.GetLayerWeightsAsTexture(
        terrainData, splatLayerIndices
      );

      string path = EditorUtility.SaveFilePanel(
        "Select Location to Save Primary Splat Map", ".", "splatmap1.png", "png"
      );
      if(!string.IsNullOrEmpty(path)) {
        writeToFile(path, splatMapTexture.EncodeToPNG());
      }
    }

    /// <summary>Saves the terrain's secondary splat map into a file</summary>
    /// <param name="terrainData">Terrain whose secondary splat map will be saved</param>
    /// <param name="splatColorIndices">Color channel assignments for the splat layers</param>
    private void exportSecondarySplatMapToTexture(
      TerrainData terrainData, int[] splatColorIndices
    ) {
      int[] splatLayerIndices = layerIndicesFromColorIndices(splatColorIndices, -5);
      Texture2D splatMapTexture = TerrainHelper.GetLayerWeightsAsTexture(
        terrainData, splatLayerIndices
      );

      string path = EditorUtility.SaveFilePanel(
        "Select Location to Save Secondary Splat Map", ".", "splatmap2.png", "png"
      );
      if(!string.IsNullOrEmpty(path)) {
        writeToFile(path, splatMapTexture.EncodeToPNG());
      }
    }

    /// <summary>Writes the contents of a byte array into a file</summary>
    /// <param name="path">Path the byte array's contents will be saved under</param>
    /// <param name="fileContents">File contents that will be saved as a file</param>
    private static void writeToFile(string path, byte[] fileContents) {
      using(
        var stream = new FileStream(path, FileMode.Create, FileAccess.Write, FileShare.None)
      ) {
        stream.Write(fileContents, 0, fileContents.Length);
        stream.Flush();
      }
    }

    /// <summary>
    ///   Converts a set of color indices per layer to a set of layer indices per color
    ///   (transposes a look up table)
    /// </summary>
    /// <param name="colorIndices">Color indices that will be converted to layer indices</param>
    /// <param name="colorOffset">Offset by which the colors will be shifted</param>
    /// <returns>Layer indices for each color</returns>
    private static int[] layerIndicesFromColorIndices(int[] colorIndices, int colorOffset = 0) {
      const int MaximumColorChannelCount = 4;

      if(colorIndices == null) {
        return null;
      }
      int colorIndexCount = colorIndices.Length;

      // Initialize the layer indices to -1 so unused layers will be clearly marked
      var layerIndices = new int[MaximumColorChannelCount];
      for(int index = 0; index < MaximumColorChannelCount; ++index) {
        layerIndices[index] = -1;
      }

      // Walk through all color channel assignments and fill the layer indices
      // at each color channel with the layer index it was found on
      for(int index = 0; index < colorIndexCount; ++index) {
        int colorChannelIndex = colorIndices[index] + colorOffset;
        if((colorChannelIndex >= 0) && (colorChannelIndex < MaximumColorChannelCount)) {
          layerIndices[colorChannelIndex] = index;
        }
      }

      return layerIndices;
    }

  }

} // namespace Framework.Terrain