using UnityEngine;
using UnityEditor;
using System.IO;

[ExecuteInEditMode]
public class Impostors : MonoBehaviour
{
    [Header("Bounding Sphere")]
    [Tooltip("Sphere radius is controlled by the Transform scale (uses X component)")]
    public float sphere_radius_ = 1f;

    [Header("Grid Settings")]
    [Tooltip("Number of lattice points along each axis (e.g., 5 = 5x5 = 25 points)")]
    [Range(2, 20)]
    public int gridResolution = 5;

    [Header("Camera Settings")]
    [Range(1, 4096)]
    public int cameraResolution = 256;

    [Tooltip("Near clipping distance - cameras are placed this distance outside the sphere")]
    public float nearClippingDistance = 0.01f;

    [Tooltip("Layers to render when baking")]
    public LayerMask cullingMask = -1;

    [Tooltip("Render skybox in baked images")]
    public bool renderSkybox = false;

    [HideInInspector]
    public GameObject[] cameraObjects;
    [HideInInspector]
    public Camera[] cameras;

    private float radius() { return sphere_radius_ * transform.lossyScale.x; }

    void OnDrawGizmos()
    {
        // Use transform scale directly for real-time gizmo updates
        float currentRadius = radius();

        // Draw the bounding sphere
        Gizmos.color = Color.cyan;
        Gizmos.DrawWireSphere(transform.position, currentRadius);

        // Draw the camera positions and directions
        if (Application.isEditor && gridResolution > 0)
        {
            Gizmos.color = Color.yellow;

            for (int y = 0; y < gridResolution; y++)
            {
                for (int x = 0; x < gridResolution; x++)
                {
                    Vector3 hemispherePos = PlaneToHemiOctahedron(x, y);
                    Vector3 worldPos = transform.position + hemispherePos * (currentRadius + nearClippingDistance);

                    // Draw camera position
                    Gizmos.DrawSphere(worldPos, currentRadius * 0.05f);

                    // Draw line to sphere center
                    Gizmos.DrawLine(worldPos, transform.position);
                }
            }
        }
    }

    // Port of plane_to_hemi_octahedron from vertex_deformation.slang
    Vector3 PlaneToHemiOctahedron(int gridX, int gridY)
    {
        // Map grid indices to [-1, 1] plane coordinates
        float u = (gridX / (float)(gridResolution - 1)) * 2f - 1f;
        float v = (gridY / (float)(gridResolution - 1)) * 2f - 1f;

        float x = u;
        float z = v;

        // Rotate 45° and scale to fit square into diamond
        float x_rot = (x + z) * 0.5f;
        float z_rot = (z - x) * 0.5f;

        // Octahedral decode: y = 1 - |x| - |z|, clamped to hemisphere
        float y = Mathf.Max(0f, 1f - Mathf.Abs(x_rot) - Mathf.Abs(z_rot));

        // Normalize to unit sphere
        Vector3 oct_pos = new Vector3(x_rot, y, z_rot);
        oct_pos.Normalize();

        // Rotate back by -45° around y to undo input rotation
        float RCP_SQRT_2 = 0.70710678f;
        float x_unrot = (oct_pos.x - oct_pos.z) * RCP_SQRT_2;
        float z_unrot = (oct_pos.x + oct_pos.z) * RCP_SQRT_2;
        oct_pos = new Vector3(x_unrot, oct_pos.y, z_unrot);

        return oct_pos;
    }

    public void CreateCameras()
    {
        // Clean up existing cameras
        DestroyExistingCameras();

        // Create parent GameObject for all cameras
        GameObject camerasParent = new GameObject("Cameras");
        camerasParent.transform.parent = transform;
        camerasParent.transform.localPosition = Vector3.zero;
        camerasParent.transform.localRotation = Quaternion.identity;
        camerasParent.transform.localScale = Vector3.one;

        int totalCameras = gridResolution * gridResolution;
        cameraObjects = new GameObject[totalCameras];
        cameras = new Camera[totalCameras];

        int index = 0;
        for (int y = 0; y < gridResolution; y++)
        {
            for (int x = 0; x < gridResolution; x++)
            {
                // Get position on hemisphere
                Vector3 hemisphereDir = PlaneToHemiOctahedron(x, y);
                float currentRadius = radius();

                // Place camera offset distance outside bounding sphere surface
                Vector3 worldPos = transform.position + hemisphereDir * (currentRadius + nearClippingDistance);

                // Create camera GameObject
                GameObject camObj = new GameObject($"ImpostorCamera_{x}_{y}");
                camObj.transform.parent = camerasParent.transform;
                camObj.transform.position = worldPos;
                camObj.transform.LookAt(transform.position);

                // Add and configure camera
                Camera cam = camObj.AddComponent<Camera>();
                cam.orthographic = true;
                cam.orthographicSize = sphere_radius_;
                cam.nearClipPlane = nearClippingDistance;
                cam.farClipPlane = sphere_radius_ * 2f + nearClippingDistance;
                cam.cullingMask = cullingMask;
                if (!renderSkybox)
                {
                    cam.clearFlags = CameraClearFlags.SolidColor;
                    cam.backgroundColor = Color.clear;
                }
                cam.enabled = false; // Only enable during baking

                cameraObjects[index] = camObj;
                cameras[index] = cam;
                index++;
            }
        }

        Debug.Log($"Created {totalCameras} impostor cameras");
    }

    public void DestroyExistingCameras()
    {
        // Find and destroy the "Cameras" parent GameObject
        Transform camerasTransform = transform.Find("Cameras");
        if (camerasTransform != null)
        {
            DestroyImmediate(camerasTransform.gameObject);
        }

        cameraObjects = null;
        cameras = null;
    }

    public void BakeTexture()
    {
        // Create a texture atlas to hold all camera views
        int texWidth = cameraResolution * gridResolution;
        int texHeight = cameraResolution * gridResolution;
        Texture2D atlasTexture = new Texture2D(texWidth, texHeight, TextureFormat.RGBA32, false);

        // Create temporary render texture for each camera
        RenderTexture rt = RenderTexture.GetTemporary(cameraResolution, cameraResolution, 24);

        int index = 0;
        for (int y = 0; y < gridResolution; y++)
        {
            for (int x = 0; x < gridResolution; x++)
            {
                Camera cam = cameras[index];

                // Render camera to RT
                cam.targetTexture = rt;
                cam.Render();

                // Read pixels from RT
                RenderTexture.active = rt;
                Texture2D temp = new Texture2D(cameraResolution, cameraResolution, TextureFormat.RGBA32, false);
                temp.ReadPixels(new Rect(0, 0, cameraResolution, cameraResolution), 0, 0);
                temp.Apply();

                // Copy to atlas at correct position
                int atlasX = x * cameraResolution;
                int atlasY = y * cameraResolution;
                atlasTexture.SetPixels(atlasX, atlasY, cameraResolution, cameraResolution, temp.GetPixels());

                DestroyImmediate(temp);
                index++;
            }
        }

        atlasTexture.Apply();
        RenderTexture.active = null;
        RenderTexture.ReleaseTemporary(rt);

        // Save texture to file
        byte[] bytes = atlasTexture.EncodeToPNG();

        // Get currently selected folder in Project window
        string folder = "Assets";
        if (Selection.activeObject != null)
        {
            string selectedPath = AssetDatabase.GetAssetPath(Selection.activeObject);
            if (!string.IsNullOrEmpty(selectedPath))
            {
                if (AssetDatabase.IsValidFolder(selectedPath))
                {
                    folder = selectedPath;
                }
                else
                {
                    folder = Path.GetDirectoryName(selectedPath);
                }
            }
        }

        string assetPath = Path.Combine(folder, "ho_bake.png");
        string fullPath = Path.Combine(Application.dataPath, "..", assetPath);
        File.WriteAllBytes(fullPath, bytes);

        Debug.Log($"Baked texture saved to: {assetPath}");
        Debug.Log($"Atlas size: {texWidth}x{texHeight} ({gridResolution}x{gridResolution} grid of {cameraResolution}x{cameraResolution} images)");

        // Refresh asset database
        AssetDatabase.Refresh();

        DestroyImmediate(atlasTexture);
    }
}

[CustomEditor(typeof(Impostors))]
public class ImpostorsEditor : Editor
{
    public override void OnInspectorGUI()
    {
        DrawDefaultInspector();

        Impostors script = (Impostors)target;

        GUILayout.Space(10);

        if (GUILayout.Button("Create Cameras", GUILayout.Height(30)))
        {
            script.CreateCameras();
        }

        if (GUILayout.Button("Destroy Cameras", GUILayout.Height(30)))
        {
            script.DestroyExistingCameras();
        }

        GUILayout.Space(10);

        bool hasCameras = script.cameras != null && script.cameras.Length > 0;

        GUI.enabled = hasCameras;
        if (GUILayout.Button(new GUIContent("Bake Texture",
            hasCameras ? "" : "Create cameras first"), GUILayout.Height(40)))
        {
            script.BakeTexture();
        }
        GUI.enabled = true;
    }
}