using System.Collections;
using System.Collections.Generic;
using Godot;
using System;
using System.Linq;

using System.Threading.Tasks;

namespace Rokojori
{
  [Tool]
  [GlobalClass]
  public partial class MeshCombiner:Node3D
  {
    [Export]
    public Node3D[] sourceNodes = [];

    [ExportToolButton( "Combine")]
    public Callable CombineButton => Callable.From( Combine );

    [ExportGroup( "Mesh")]
    [Export]
    public bool combineMeshes = true;
    
    public enum UVCombineMode
    {
      Keep,
      Adjust_To_Combined_Material
    }

    [Export]
    public UVCombineMode uVCombineMode = UVCombineMode.Adjust_To_Combined_Material;

    [Export]
    public Node3D pivot;

    [ExportGroup( "Material")]
    [Export]
    public bool combineMaterials = true; 
    
    [Export]
    public bool combineORM = true;

    public enum TextureSizeMode
    {
      KeepOriginal,
      Custom
    }

    [Export]
    public TextureSizeMode textureSizeMode = TextureSizeMode.KeepOriginal;

    [Export]
    public Vector2I customTextureSize = new Vector2I( 1024, 1024 );
    
    
    [ExportGroup( "Output")]
    [Export]
    public MeshInstance3D outputMesh;

    [Export]
    public Node3D outputContainer;

    [Export]
    public Material[] outputMaterials = [];

    

    MultiMap<Mesh,int,MeshGeometry> _meshGeometryCache = new MultiMap<Mesh,int,MeshGeometry>();

    MeshGeometry GetMeshGeometry( MeshSurface surface )
    {
      if ( ! _meshGeometryCache.Has( surface.mesh, surface.index ) )
      {
        var mg = MeshGeometry.From( surface.mesh as ArrayMesh, null, surface.index );
        _meshGeometryCache.Set( surface.mesh, surface.index, mg );

        this.LogInfo( "Created mesh with triangles:", mg.numTriangles );
      }

      return _meshGeometryCache[ surface.mesh ][ surface.index ];
    }
    
    List<Transformable<MeshSurface>> _surfaces = new List<Transformable<MeshSurface>>();
    List<Material> _materials = new List<Material>();
    MapList<Material,Transformable<MeshSurface>> _materiaList = new MapList<Material, Transformable<MeshSurface>>() ;
    Dictionary<Material,Transform2D> _uvTransform = new Dictionary<Material, Transform2D>();
    MapList<MaterialType,Material> _materialTypeGroups = new MapList<MaterialType, Material>();
    Dictionary<Material,Material> _combinedMaterials = new Dictionary<Material, Material>();

    void Clear()
    {
      _surfaces.Clear();
      _materials.Clear();
      _materiaList.Clear();
      _uvTransform.Clear();
      _materialTypeGroups.Clear();
      _combinedMaterials.Clear();      

      outputMesh = null;
    }

    public async Task Combine()
    {
      Clear();

      try
      {
        await GrabSurfaces();
        await GrabMaterials();
        
        await CombineMaterials();

        await CombineMeshes();
        
        
        
      }
      catch ( System.Exception e )
      {
        this.LogError( e );
      }

      return;
    }

    async Task GrabSurfaces()
    {     
      _surfaces = new List<Transformable<MeshSurface>>();
      
      foreach ( var n in sourceNodes )
      {      
        MeshExtractor.ExtractSurfacesInHierarchy( n, _surfaces ); 
      }
    }

    async Task GrabMaterials()
    {
      _materials = new List<Material>();
      _materiaList = new MapList<Material, Transformable<MeshSurface>>();

      var set = new HashSet<Material>();

      _surfaces.ForEach(
        ( s )=>
        {
          _materiaList.Add( s.item.material, s );
          
          if ( set.Contains( s.item.material ) )
          {
            return;
          }

          set.Add( s.item.material );
          _materials.Add( s.item.material );
          
        }
      );


    }

    async Task CombineMaterials()
    {
      _materials.ForEach(
        m =>
        {
          var type = _materialTypeGroups.FindKey( k => k.EqualsTo( m ) );
          
          if ( type == null )
          {
            type = MaterialType.From( m );

            if ( type == null )
            {
              this.LogInfo( "Invalid Type:", m);
            }
            
          }

          if ( type == null )
          {
            return;
          }

          _materialTypeGroups.Add( type, m );
        }
      );

      this.LogInfo( "Found Groups:", _materialTypeGroups.Count );

      foreach ( var gm in _materialTypeGroups )
      { 
        this.LogInfo( "Group:", gm.Key, ">>", gm.Value.Count );

        if ( gm.Value.Count > 1 )
        {
          await CombineStandardMaterials( gm.Value );
        }        
        
      }

    }

    async Task<StandardMaterial3D> CombineStandardMaterials( List<Material> anyMaterials )
    {
      var materials = Lists.FilterType<Material,StandardMaterial3D>( anyMaterials );

      var outputMaterial = (StandardMaterial3D) materials[ 0 ].Duplicate();

      var alignment = TextureMerger.ComputeTextureAlignment( anyMaterials.Count );

      for ( int i = 0; i < materials.Count; i++ )
      {
        var box = TextureMerger.GetUVAlignmentBoxFor( alignment, i );
        var transform = Transform2D.Identity;
        transform = transform.ScaledLocal( box.size );
        transform.Origin = box.min;

        _uvTransform[ materials[ i ] ] = transform;
        _combinedMaterials[ materials[ i ] ] = outputMaterial;

        this.LogInfo( "Set UV transform:", i, box.min, box.size.Inverse() );
      }

      var textureTypes = new List<string>
      {
        "albedo",
        "metallic",
        "roughness",
        "emission",
        "normal",
        "rim",
        "clearcoat",
        "anisotropy_flowmap",
        "ao",
        "height",
        "subsurf_scatter",
        "backlight",
        "refraction"
      };

      var hasAlebdoAlpha = true;
      
      var mm = new StandardMaterial3D();

      foreach ( var t in textureTypes )
      {
        if ( combineORM && ( t == "roughness" || t == "metallic" ) )
        {
          continue;
        }

        var hasAlpha = hasAlebdoAlpha && t == "albedo";

        var name = Texture2DPropertyName.Create( t + "_texture" ); 
        var textures = Lists.Map( materials, m => name.Get( m ) );

        var noTextures = textures.Find( t => t != null ) == null;

        if ( noTextures )
        {
          continue;
        }

        var fallBackColors = Lists.Map(
          textures, i => 
          {
            if ( t == "ao" && combineORM )
            {
              return new Color( 1, 1, 0 );
            } 

            if ( t == "albedo" )
            {
              return new Color( 1, 1, 1 );
            } 

            if ( t == "normal" )
            {
              return new Color( 0.5f, 0.5f, 1 );
            }

            return new Color( 0, 0, 0 );
          }
        );
 
        

        var textureSize = new Vector2( 512, 512 );

        if ( TextureSizeMode.Custom == textureSizeMode )
        {
          textureSize = customTextureSize;
        }
        else if ( TextureSizeMode.KeepOriginal == textureSizeMode )
        {
          textures.ForEach(
            t =>
            {
              if ( t == null )
              {
                return;
              }

              textureSize.X = Mathf.Max( t.GetWidth(),textureSize.X );
              textureSize.Y = Mathf.Max( t.GetWidth(),textureSize.Y );
            }
          );

          this.LogInfo( "Computing next power of two texture Size", textureSize, alignment );

          textureSize.X = MathX.NextPowerOfTwo( textureSize.X * alignment.X );
          textureSize.Y = MathX.NextPowerOfTwo( textureSize.Y * alignment.Y );
        }

        this.LogInfo( "Texture Size", t, ">>", textures.Count, textureSize, alignment );


        var texture = TextureCombinerBuffer.GridMerge( 
          (int) textureSize.X, (int)textureSize.Y, 
          new Color( 1, 0, 0, hasAlpha ? 0 : 1 ), 
          hasAlpha, true, alignment,
          textures, fallBackColors
        );

        this.LogInfo( "Combined textures:", t );

        name.Set( outputMaterial, texture );
        
     
        await this.RequestNextFrame();

      }


      return outputMaterial;
    }

    async Task CombineMeshes()
    {
      var arrayMesh = new ArrayMesh();

      this.LogInfo( "Combining", _surfaces.Count, "meshes" );

      var index = 0;

      var combined = new MapList<Material,MeshGeometry>();

      var outputMaterials = new List<Material>();
      
      if ( ! combineMeshes )
      {
        if ( outputContainer == null )
        {
          outputContainer = this.CreateChild<Node3D>();
        }

        outputContainer.DestroyChildren();
      }

      _materials.ForEach(
        ( m )=>
        {
          var isCombinedMaterial = _combinedMaterials.ContainsKey( m );
          var usedMaterial = isCombinedMaterial ? _combinedMaterials[ m ] : m;

          Transform2D? uvTransform = _uvTransform.ContainsKey( m ) ? _uvTransform[ m ] : null;

          var surfaces = _materiaList[ m ];

          this.LogInfo( "Combining for Material", "combined?:",isCombinedMaterial,"material:", usedMaterial, surfaces.Count, "meshes" );

          var meshGeometry = new MeshGeometry();

          surfaces.ForEach(
            ( s )=>
            {
              var smg = GetMeshGeometry( s.item ).Clone();

              if ( uvTransform != null )
              {
                this.LogInfo( "Appling uvTransform:", uvTransform );
                smg.ApplyUVTransform( (Transform2D) uvTransform );
              }

              var trsf = s.transform;

              if ( combineMeshes )
              {
                if ( pivot != null )
                {
                  trsf.Origin -= pivot.GlobalPosition;                
                }         

                smg.ApplyTransform( trsf );     
              }
                

              meshGeometry.Add( smg );

              if ( isCombinedMaterial && ! combineMeshes )
              {
                meshGeometry.GenerateMesh( Mesh.PrimitiveType.Triangles, arrayMesh );
                arrayMesh.SurfaceSetMaterial( index, usedMaterial );        

                var meshInstance = outputContainer.CreateChild<MeshInstance3D>();
                meshInstance.Mesh = arrayMesh;

                meshInstance.GlobalTransform = trsf;

                meshGeometry = new MeshGeometry();    
                arrayMesh = new ArrayMesh();
              }
            }
          );

          

          if ( isCombinedMaterial && combineMeshes )
          {
            this.LogInfo( "Add material groups", m );
            combined.Add( _combinedMaterials[ m ], meshGeometry );
            return;
          }

          
          outputMaterials.Add( usedMaterial );

          if ( combineMeshes )
          {
            meshGeometry.GenerateMesh( Mesh.PrimitiveType.Triangles, arrayMesh );
            arrayMesh.SurfaceSetMaterial( index, usedMaterial );            

            index ++;
          }

        }
      );

      this.LogInfo( "Combining material groups", combined.Count );

      foreach ( var cm in combined )
      {
        var material = cm.Key;
        var meshes = cm.Value;
        this.LogInfo( "Combining meshes", meshes.Count);
        var combinedMG = MeshGeometry.Combine( meshes );

        this.LogInfo( "Combed meshes, num tris:", combinedMG.numTriangles );
        combinedMG.GenerateMesh( Mesh.PrimitiveType.Triangles, arrayMesh );

        this.LogInfo( "Add surface", index);
        arrayMesh.SurfaceSetMaterial( index, material );

        this.LogInfo( "Add material", material );
        outputMaterials.Add( material );

        this.LogInfo( "Increment index" );
        index ++;
      }

      this.LogInfo( "Processing done, adding outputs", arrayMesh.GetSurfaceCount() );

      this.outputMaterials = outputMaterials.ToArray();

      if ( combineMeshes )
      {       
        if ( outputMesh == null )
        {
          this.LogInfo( "Created outputMesh");
          outputMesh = this.CreateChild<MeshInstance3D>();
        }       

        outputMesh.Mesh = arrayMesh;

      }

      return;
    } 

    
  }
}