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

using Rokojori.Extensions;
namespace Rokojori
{
  /** Renders a height map from the HeightMapDrawLayer nodes underneath it.

      The node is the map's placement in the world: its global XZ is the map's
      center, `size` is how many meters the map covers, and `resolution` is how
      many texels that is. So a layer is positioned by <marked>moving it in the
      3D viewport</marked> - the generator converts world XZ to texels, never the
      other way round, and nothing in the system has a texel-space authoring
      concept at all.

      Height is stored normalized in the R channel and decoded by
      addTerrainOffset ( Shading/Library/Terrain.gdshaderinc ) as
      mix( minHeight, maxHeight, r ). minHeight/maxHeight are therefore part of
      the contract between this node and the terrain material, which is what
      PushToMaterials() exists to keep in sync.

      Rendering is flag-driven, not immediate: a layer that moves or changes
      calls SetDirty(), and the next _Process re-renders once. Several edits in
      one frame - dragging a node, which moves it and re-emits every frame -
      collapse into a single render. The render itself is GPU blits only, so at
      the sizes this is meant for it is affordable per frame, which is what makes
      dragging a layer around feel live in the editor.

      Layers are collected by walking the children, so <marked>nesting is
      grouping</marked>: put layers under a plain Node3D to move them together.
      A nested HeightMapGenerator stops the walk - it owns its own map.
  */
  [Tool]
  [GlobalClass]
  public partial class HeightMapGenerator:Node3D
  {
    static readonly Color White = new Color( 1, 1, 1, 1 );
    static readonly Color Transparent = new Color( 0, 0, 0, 0 );

    [ExportToolButton( "Render" )]
    public Callable RenderButton => Callable.From( () => Render() );


    Vector2 _size = new Vector2( 100, 100 );

    /** The map's extent in world meters, on X and Z. */
    [Export]
    public Vector2 size
    {
      get => _size;
      set { _size = value; SetDirty(); }
    }

    Vector2I _resolution = new Vector2I( 512, 512 );

    /** The map's texture size. Combined with `size` this is the texel density -
        512 over 100 m is roughly 0.2 m per texel.
    */
    [Export]
    public Vector2I resolution
    {
      get => _resolution;
      set { _resolution = value; SetDirty(); }
    }

    float _minHeight = 0;

    /** World height a stored 0 decodes to. Must match terrainMinHeight on the
        material - PushToMaterials() does that.
    */
    [Export]
    public float minHeight
    {
      get => _minHeight;
      set { _minHeight = value; SetDirty(); }
    }

    float _maxHeight = 100;

    /** World height a stored 1 decodes to. Everything a layer writes is a
        fraction of maxHeight - minHeight, so widening the range flattens the
        existing layers rather than rescaling them.
    */
    [Export]
    public float maxHeight
    {
      get => _maxHeight;
      set { _maxHeight = value; SetDirty(); }
    }

    float _baseHeight = 0;

    /** The world height the map starts at before any layer draws. */
    [Export]
    public float baseHeight
    {
      get => _baseHeight;
      set { _baseHeight = value; SetDirty(); }
    }

    TextureStackFormat _format = TextureStackFormat.RgbaHalf;

    /** RgbaHalf by default, not Rgba8: 8 bits give 256 height steps over the
        whole min-to-max range, which terraces visibly on any slope. Half is
        smooth and still only 8 bytes per texel.
    */
    [Export]
    public TextureStackFormat format
    {
      get => _format;
      set { _format = value; SetDirty(); }
    }

    /** The map itself. A DrawableTexture2D because that is what can be blitted
        into; it is still a plain Texture2D to everything that samples it, so it
        goes straight into a material's sampler slot.

        Left empty, one is created on the first render - but then only this node
        holds it, so the material has to be fed through `materials` below. Assign
        a shared DrawableTexture2D here to reference the same map from a material
        in the scene file instead.
    */
    [Export]
    public DrawableTexture2D heightMapTarget;

    /** World-space normals baked out of the finished height map, encoded as
        n * 0.5 + 0.5 - see HeightMapNormals.gdshader.

        Left empty, no normal pass runs at all and a terrain material has to
        take its own derivative per fragment. Assigned, the map is rendered once
        after every layer has drawn and pushed to `materials` as
        terrainNormalMap, which is both cheaper per fragment and the only way the
        smoothing below can exist - a consumer taking its own 1-texel central
        difference has no way to be smoothed.
    */
    [Export]
    public DrawableTexture2D normalMapTarget;

    float _normalRadius = 1;

    /** How far apart the normal pass's taps are, in texels. Widening the stencil
        low-passes the slope, which is what removes texel stair-stepping without
        touching the height map itself.
    */
    [Export]
    public float normalRadius
    {
      get => _normalRadius;
      set { _normalRadius = value; SetDirty(); }
    }

    int _normalSamples = 1;

    /** How many concentric stencils, at 1x / 2x / ... normalRadius, get
        averaged. One wide stencil misses the detail between its taps; averaging
        several keeps the fine slope and still drops the noise. Costs 8 taps
        each, over the whole map.
    */
    [Export( PropertyHint.Range, "1,8" )]
    public int normalSamples
    {
      get => _normalSamples;
      set { _normalSamples = value; SetDirty(); }
    }

    float _normalStrength = 1;

    /** Scales the slope before it is normalized - below 1 flattens the shading,
        above 1 exaggerates it. Shading only; the geometry is unaffected.
    */
    [Export]
    public float normalStrength
    {
      get => _normalStrength;
      set { _normalStrength = value; SetDirty(); }
    }

    /** Materials to push terrainHeightMap / terrainNormalMap / terrainCenterXZ / terrainSizeXZ /
        terrainMinHeight / terrainMaxHeight onto after every render.

        This is the whole wiring between the generator and anything rendering the
        terrain - GeometryTerrainOffset's uniform block, by name. Without it the
        material would silently keep the size and height range it was authored
        with while the generator used different ones.
    */
    [Export]
    public Material[] materials = new Material[ 0 ];

    /** Off makes the map update only on the Render button, for a map too big to
        recompute while dragging.
    */
    [Export]
    public bool autoRender = true;

    [ExportGroup( "Debug Info" )]

    [Export]
    public int renderedLayers;


    /* THE DIRTY FLAG
       ------------------------------------------------------------------------
       One bool, checked in _Process, rather than rendering straight out of the
       setter. A drag emits a transform notification every frame and an inspector
       slider emits one per pixel of mouse travel, so rendering per notification
       would run several full renders per frame to show one result.
    */

    bool _dirty = true;

    public void SetDirty()
    {
      _dirty = true;
    }

    public bool isDirty => _dirty;

    public override void _Ready()
    {
      SetDirty();
    }

    public override void _Process( double delta )
    {
      if ( ! _dirty || ! autoRender )
      {
        return;
      }

      Render();
    }


    /* COORDINATES
       ------------------------------------------------------------------------
       Texel (0,0) is the map's minimum XZ corner. addTerrainOffset computes its
       UV as ( worldXZ - ( center - size/2 ) ) / size and samples with texture(),
       whose (0,0) is the image's first texel - and BlitRect's Rect2I origin is
       that same first texel, with no Y flip. So image +Y is world +Z throughout,
       and nothing here needs to mirror anything.
    */

    /** World XZ of the map's center - the node's own global position. */
    public Vector2 centerXZ
    {
      get
      {
        var p = GlobalPosition;
        return new Vector2( p.X, p.Z );
      }
    }

    public Vector2 minXZ => centerXZ - size / 2;

    public Vector2I textureSize => new Vector2I( Mathf.Max( 1, resolution.X ), Mathf.Max( 1, resolution.Y ) );

    public Vector2 metersPerTexel
    {
      get
      {
        var r = textureSize;
        return new Vector2( size.X / r.X, size.Y / r.Y );
      }
    }

    public Rect2I fullRect => new Rect2I( 0, 0, textureSize.X, textureSize.Y );

    public Vector2 TexelToWorldXZ( Vector2 texel )
    {
      var m = metersPerTexel;
      var origin = minXZ;

      return new Vector2( origin.X + texel.X * m.X, origin.Y + texel.Y * m.Y );
    }

    public Vector2 WorldXZToTexel( Vector2 worldXZ )
    {
      var m = metersPerTexel;
      var origin = minXZ;

      return new Vector2( ( worldXZ.X - origin.X ) / m.X, ( worldXZ.Y - origin.Y ) / m.Y );
    }

    /** The texel rect a world-space XZ box covers, grown to whole texels and
        clipped to the map. Empty when the box misses the map entirely.
    */
    public Rect2I WorldBoxToRect( Vector2 worldMin, Vector2 worldMax )
    {
      var min = WorldXZToTexel( worldMin );
      var max = WorldXZToTexel( worldMax );

      var x0 = Mathf.FloorToInt( Mathf.Min( min.X, max.X ) );
      var y0 = Mathf.FloorToInt( Mathf.Min( min.Y, max.Y ) );
      var x1 = Mathf.CeilToInt( Mathf.Max( min.X, max.X ) );
      var y1 = Mathf.CeilToInt( Mathf.Max( min.Y, max.Y ) );

      return new Rect2I( x0, y0, x1 - x0, y1 - y0 ).Intersection( fullRect );
    }

    /** Convenience for a radially symmetric layer. */
    public Rect2I WorldCircleToRect( Vector3 worldCenter, float worldRadius )
    {
      var c = new Vector2( worldCenter.X, worldCenter.Z );
      var r = new Vector2( worldRadius, worldRadius );

      return WorldBoxToRect( c - r, c + r );
    }

    /** The texel rect a rotated XZ quad covers - the axis-aligned bounds of its
        four corners. A rotated layer therefore touches more texels than it
        fills; the shader still rejects everything outside the quad itself.
    */
    public Rect2I WorldQuadToRect( Vector2 center, Vector2 axisX, Vector2 axisZ, Vector2 size )
    {
      var x = axisX * ( size.X / 2 );
      var z = axisZ * ( size.Y / 2 );

      var a = center + x + z;
      var b = center + x - z;
      var c = center - x + z;
      var d = center - x - z;

      var min = new Vector2(
        Mathf.Min( Mathf.Min( a.X, b.X ), Mathf.Min( c.X, d.X ) ),
        Mathf.Min( Mathf.Min( a.Y, b.Y ), Mathf.Min( c.Y, d.Y ) )
      );

      var max = new Vector2(
        Mathf.Max( Mathf.Max( a.X, b.X ), Mathf.Max( c.X, d.X ) ),
        Mathf.Max( Mathf.Max( a.Y, b.Y ), Mathf.Max( c.Y, d.Y ) )
      );

      return WorldBoxToRect( min, max );
    }


    /* HEIGHT NORMALIZATION
       ------------------------------------------------------------------------
       Layers author in meters; the map stores 0..1. These two are the only
       conversion, and which one a layer uses is decided by its blend mode -
       HeightMapBlendMode.IsAbsolute().
    */

    public float heightRange
    {
      get
      {
        var range = maxHeight - minHeight;

        return Mathf.Abs( range ) < 0.0001f ? 1 : range;
      }
    }

    /** A height difference in meters as a fraction of the range, for Add and
        Subtract.
    */
    public float NormalizeDelta( float worldHeightDelta )
    {
      return worldHeightDelta / heightRange;
    }

    /** An absolute world height as a stored 0..1 value, for Replace. */
    public float NormalizeLevel( float worldHeight )
    {
      return ( worldHeight - minHeight ) / heightRange;
    }

    public float baseLevel => NormalizeLevel( baseHeight );


    /* RENDERING */

    HeightMapShapeMaterial _shapeMaterial;

    /** The shared content-pass material. Layers configure it in
        RenderToHeightMap() and then call DrawShape(), rather than allocating a
        material each.
    */
    public HeightMapShapeMaterial shapeMaterial
    {
      get
      {
        if ( _shapeMaterial == null )
        {
          _shapeMaterial = new HeightMapShapeMaterial();
        }

        return _shapeMaterial;
      }
    }

    BlitMaterial _blendBlit;
    BlitMaterial _copyBlit;

    /** Straight overwrite, for clearing. */
    BlitMaterial copyBlit
    {
      get
      {
        if ( _copyBlit == null )
        {
          _copyBlit = new BlitMaterial();
          _copyBlit.BlendMode = BlitMaterial.BlendModeEnum.Disabled;
        }

        return _copyBlit;
      }
    }

    readonly List<HeightMapDrawLayer> _layers = new List<HeightMapDrawLayer>();

    bool _rendering = false;

    /** Clears the map to baseHeight and lets every visible layer underneath draw
        into it, in tree order - so the layer listed last wins, the same way the
        3D viewport draws it.
    */
    public void Render()
    {
      // A layer that calls SetDirty() from inside its own draw - a transform
      // notification from a tool script, say - would otherwise recurse.
      if ( _rendering )
      {
        return;
      }

      _dirty = false;
      _rendering = true;

      try
      {
        PrepareTarget();

        _layers.Clear();
        CollectLayers( this, _layers );

        renderedLayers = 0;

        for ( var i = 0; i < _layers.Count; i++ )
        {
          var layer = _layers[ i ];

          if ( ! layer.IsVisibleInTree() )
          {
            continue;
          }

          layer.RenderToHeightMap( this );

          renderedLayers++;
        }

        RenderNormalMap();

        PushToMaterials();
      }
      finally
      {
        _rendering = false;
      }
    }

    void CollectLayers( Node node, List<HeightMapDrawLayer> into )
    {
      foreach ( var child in node.GetChildren() )
      {
        // Its own map, its own layers.
        if ( child is HeightMapGenerator )
        {
          continue;
        }

        if ( child is HeightMapDrawLayer layer )
        {
          into.Add( layer );
        }

        CollectLayers( child, into );
      }
    }

    /** One blit per layer pair: the shape into a rect-sized scratch buffer with
        blending disabled, then the scratch onto the map with the blend mode's
        hardware blend.

        The scratch buffer is rect-sized rather than full-size because BlitRect
        stretches its whole source across the destination rect - a full-size
        source would be squeezed into the rect.

        Called by layers after they have set up `shapeMaterial`; the rect and
        blend-mode uniforms are set here, so a layer cannot get the world mapping
        wrong.

        `source` and `mask` fill the shader's hint_blit_source0 and
        hint_blit_source1 - an image layer passes its texture and optionally a
        mask, a purely procedural shape passes neither. Both slots are always
        bound, to the shared white 1x1 stand-in when a layer has nothing for
        them: an unbound blit source arrives black, so "no mask" would otherwise
        have to be a flag, while white is already a no-op multiplier.
    */
    public void DrawShape( Rect2I rect, HeightMapBlendMode blendMode )
    {
      DrawShape( rect, blendMode, null, null );
    }

    public void DrawShape( Rect2I rect, HeightMapBlendMode blendMode, Texture2D source )
    {
      DrawShape( rect, blendMode, source, null );
    }

    public void DrawShape( Rect2I rect, HeightMapBlendMode blendMode, Texture2D source, Texture2D mask )
    {
      if ( heightMapTarget == null || rect.Size.X <= 0 || rect.Size.Y <= 0 )
      {
        return;
      }

      var material = shapeMaterial;
      var meters = metersPerTexel;

      material.rectWorldMin.Set( TexelToWorldXZ( new Vector2( rect.Position.X, rect.Position.Y ) ) );
      material.rectWorldSize.Set( new Vector2( rect.Size.X * meters.X, rect.Size.Y * meters.Y ) );
      material.blendMode.Set( (int) blendMode );

      var scratch = TextureStackBufferPool.shared.Rent( rect.Size.X, rect.Size.Y, format, Transparent );

      try
      {
        _blitSources[ 0 ] = source != null ? source : (Texture2D) dummySource;
        _blitSources[ 1 ] = mask != null ? mask : (Texture2D) dummySource;

        // BlitRectMulti rather than BlitRect, because the explicit sources array
        // is the documented way to fill more than source 0.
        scratch.BlitRectMulti(
          new Rect2I( 0, 0, rect.Size.X, rect.Size.Y ),
          _blitSources, TextureStackCompositor.noExtraTargets, White, 0, material
        );

        if ( _blendBlit == null )
        {
          _blendBlit = new BlitMaterial();
        }

        _blendBlit.BlendMode = blendMode.ToBlitBlendMode();

        heightMapTarget.BlitRect( rect, scratch, White, 0, _blendBlit );
      }
      finally
      {
        TextureStackBufferPool.shared.Return( scratch );
      }
    }

    /** Reused rather than allocated per blit - a drag re-renders every frame. */
    readonly Godot.Collections.Array<Texture2D> _blitSources
      = new Godot.Collections.Array<Texture2D>(){ null, null };

    DrawableTexture2D _dummySource;

    DrawableTexture2D dummySource
    {
      get
      {
        if ( _dummySource == null )
        {
          _dummySource = new DrawableTexture2D();
          _dummySource.Setup( 1, 1, DrawableTexture2D.DrawableFormat.Rgba8, White, false );
        }

        return _dummySource;
      }
    }


    /* PREPARING THE MAP
       ------------------------------------------------------------------------
       Setup() reallocates, and reallocating makes the engine emit the texture's
       changed signal - which rebuilds the inspector. So it runs only when the
       allocation actually has to change, and every other render clears by
       blitting a 1x1 base-height texture across the map with blending disabled.
       Same pixels, no reallocation, no signal. This is also what takes the one
       expensive fixed cost out of a re-render and makes rendering per frame
       affordable - the same finding the TextureStack compositor landed on.
    */

    DrawableTexture2D _setupTarget;
    Vector2I _setupSize;
    TextureStackFormat _setupFormat;

    DrawableTexture2D _clearSource;
    Color _clearSourceColor;
    TextureStackFormat _clearSourceFormat;

    void PrepareTarget()
    {
      if ( heightMapTarget == null )
      {
        heightMapTarget = new DrawableTexture2D();
      }

      var level = baseLevel;
      var clearColor = new Color( level, level, level, 1 );
      var wanted = textureSize;

      if ( _setupTarget != heightMapTarget || _setupSize != wanted || _setupFormat != format )
      {
        heightMapTarget.Setup( wanted.X, wanted.Y, format.ToDrawableFormat(), clearColor, false );

        _setupTarget = heightMapTarget;
        _setupSize = wanted;
        _setupFormat = format;

        return;
      }

      ClearTarget( clearColor );
    }

    void ClearTarget( Color clearColor )
    {
      if ( _clearSource == null || _clearSourceColor != clearColor || _clearSourceFormat != format )
      {
        if ( _clearSource == null )
        {
          _clearSource = new DrawableTexture2D();
        }

        _clearSource.Setup( 1, 1, format.ToDrawableFormat(), clearColor, false );

        _clearSourceColor = clearColor;
        _clearSourceFormat = format;
      }

      heightMapTarget.BlitRect( fullRect, _clearSource, White, 0, copyBlit );
    }


    /* THE NORMAL PASS
       ------------------------------------------------------------------------
       One full-map blit with the finished height map as the source. It runs
       after the layers rather than per layer, because a normal is a derivative
       of the composited result - taking it per layer and blending the results
       would be wrong wherever two layers overlap.

       Nothing clears the normal map: the shader writes every texel with blending
       disabled, so the previous render is fully overwritten.
    */

    HeightMapNormalsMaterial _normalsMaterial;

    HeightMapNormalsMaterial normalsMaterial
    {
      get
      {
        if ( _normalsMaterial == null )
        {
          _normalsMaterial = new HeightMapNormalsMaterial();
        }

        return _normalsMaterial;
      }
    }

    DrawableTexture2D _setupNormalTarget;
    Vector2I _setupNormalSize;
    TextureStackFormat _setupNormalFormat;

    /** Flat up, encoded - what an unrendered normal map should read as. */
    static readonly Color FlatNormal = new Color( 0.5f, 1, 0.5f, 1 );

    void RenderNormalMap()
    {
      if ( normalMapTarget == null || heightMapTarget == null )
      {
        return;
      }

      var wanted = textureSize;

      // Same allocation rule as the height map - Setup() reallocates and that
      // emits the texture's changed signal, which rebuilds the inspector.
      if ( _setupNormalTarget != normalMapTarget
        || _setupNormalSize != wanted
        || _setupNormalFormat != format )
      {
        normalMapTarget.Setup( wanted.X, wanted.Y, format.ToDrawableFormat(), FlatNormal, false );

        _setupNormalTarget = normalMapTarget;
        _setupNormalSize = wanted;
        _setupNormalFormat = format;
      }

      var material = normalsMaterial;

      material.sizeXZ.Set( size );
      material.minHeight.Set( minHeight );
      material.maxHeight.Set( maxHeight );
      material.radius.Set( normalRadius );
      material.samples.Set( normalSamples );
      material.strength.Set( normalStrength );

      normalMapTarget.BlitRect( fullRect, heightMapTarget, White, 0, material );
    }


    /** Pushes the maps and their world mapping onto every assigned material. */
    public void PushToMaterials()
    {
      if ( materials == null )
      {
        return;
      }

      for ( var i = 0; i < materials.Length; i++ )
      {
        var material = materials[ i ];

        if ( material == null )
        {
          continue;
        }

        TerrainOffsetShader.heightMap.Set( material, heightMapTarget );
        TerrainOffsetShader.normalMap.Set( material, normalMapTarget );
        TerrainOffsetShader.centerXZ.Set( material, centerXZ );
        TerrainOffsetShader.sizeXZ.Set( material, size );
        TerrainOffsetShader.minHeight.Set( material, minHeight );
        TerrainOffsetShader.maxHeight.Set( material, maxHeight );
      }
    }


    /** Rough VRAM cost of the maps, in bytes. The scratch buffer is rect-sized
        and pooled, so it is not counted - a layer covering the whole map would
        add one more of these.
    */
    public long EstimateMemory()
    {
      var r = textureSize;
      var maps = normalMapTarget != null ? 2 : 1;

      return (long) r.X * r.Y * format.BytesPerPixel() * maps;
    }
  }
}