using UnityEngine;
using System.Collections.Generic;
using UnityEditor;
using System;

public class ProceduralGrass : MonoBehaviour
{
    [Header("References")]
    [SerializeField] private Mesh grassMesh;
    [SerializeField] private Material grassMaterial;
    [SerializeField] private GrassPointsData grassPointsData;

    [Header("Grass Generation: ")]
    [SerializeField] private int rejectionSamples = 12000;
    [SerializeField] private float radius = 0.1f;

    [Header("Grass Settings")]
    [SerializeField] private Vector2 regionSize = new Vector2(30f, 30f);
    [SerializeField] private Vector2 grassHeightRange = new Vector2 (0, 0);
    [SerializeField] private Vector2 grassYRotationRange = new Vector2 (0, 0);

    private int vertexCount;

    public Vector2 GrassHeightRange { get => grassHeightRange; set => grassHeightRange = value; }
    public Vector2 GrassYRotationRange { get => grassYRotationRange; set => grassYRotationRange = value; }

#if UNITY_EDITOR
[ContextMenu("Generate and Save Points")]
    void GenerateAndSavePoints()
    {
        if (grassMesh == null || grassMaterial == null || grassPointsData == null)
        {
            Debug.LogError("Missing references in ProceduralGrass!");
            return;
        }

        List<Vector2> points = PoissonDiscSamplingJobSystem.GeneratePoints(radius, regionSize, rejectionSamples);
        grassPointsData.points.Clear();
        grassPointsData.points.AddRange(points);

        EditorUtility.SetDirty(grassPointsData);
        AssetDatabase.SaveAssets();
        Debug.Log($"Saved {points.Count} grass points to {grassPointsData.name} using Job System + Burst");
    }
#endif

    struct GrassBladeData
    {
        public Matrix4x4 transform;

        public GrassBladeData(Vector3 position, float rotationY, float scale)
        {
            Quaternion rotation = Quaternion.Euler(0, rotationY, 0); //include Rotation
            Vector3 scaleVec = Vector3.one * scale; //Calculate and include the scale
            transform = Matrix4x4.TRS(position, rotation, scaleVec); //include the position
        }
    }

    private GraphicsBuffer argsBuffer;
    private GraphicsBuffer instanceDataBuffer;
    private MaterialPropertyBlock materialProps;

    void Start()
    {
        // Initialize materialProps FIRST
        materialProps = new MaterialPropertyBlock();

        if (grassPointsData == null || grassPointsData.points.Count == 0)
        {
            Debug.LogError("No grass points data!");
            return;
        }

        // Fill grass blades
        GrassBladeData[] blades = new GrassBladeData[grassPointsData.points.Count];
        Matrix4x4[] bladeMatrices = new Matrix4x4[blades.Length];

        int numberOfIterations = 0;
        for (int i = 0; i < blades.Length; i++)
        {
            Vector2 point = grassPointsData.points[i];

            //Hardcoded Y position
            Vector3 pos = new Vector3(point.x - regionSize.x / 2f, -0.01f, point.y - regionSize.y / 2f);
            float scale = 0.05f * UnityEngine.Random.Range(grassHeightRange.x, grassHeightRange.y);
            float rotation = UnityEngine.Random.Range(grassYRotationRange.x, grassYRotationRange.y);

            blades[i] = new GrassBladeData(pos, rotation, scale);
            bladeMatrices[i] = blades[i].transform;
            numberOfIterations++;
        }

        Debug.Log("Number of iterations: " + numberOfIterations);

        // Setup instance data buffer
        instanceDataBuffer = new GraphicsBuffer(GraphicsBuffer.Target.Structured, bladeMatrices.Length, sizeof(float) * 16);
        instanceDataBuffer.SetData(bladeMatrices);
        Debug.Log("First blade matrix: " + bladeMatrices[0]);

        // Now materialProps is initialized
        materialProps.SetBuffer("_GrassBladeData", instanceDataBuffer); // Make sure this matches shader

        //Get Vertex number
        if (grassMesh != null)
        {
            vertexCount = grassMesh.vertexCount;

            float v = Convert.ToSingle(vertexCount);
            materialProps.SetFloat("_Mesh_Vertex_Number", v);
        }

        // Setup args buffer
        argsBuffer = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, 1, sizeof(uint) * 5);
        uint[] args = new uint[5]
        {
        grassMesh.GetIndexCount(0),
        (uint)blades.Length,
        grassMesh.GetIndexStart(0),
        grassMesh.GetBaseVertex(0),
        0
        };
        argsBuffer.SetData(args);
    }

    void Update()
    {
        if (!Application.isPlaying)
            return;

        Graphics.RenderMeshIndirect(
            new RenderParams(grassMaterial)
            {
                worldBounds = new Bounds(transform.position, new Vector3(regionSize.x, 100, regionSize.y)),
                matProps = materialProps
            },
            grassMesh,
            argsBuffer
        );
    }

    void OnDestroy()
    {
        argsBuffer?.Release();
        instanceDataBuffer?.Release();
        instanceDataBuffer = null;
    }

}