using UnityEngine;
using UnityEditor;
using System.IO;

[ExecuteInEditMode]
public class Impostors : MonoBehaviour
{
  [Header("Bounding Sphere")]
  public float sphere_radius_ = 1f;

  [Header("Grid Settings")]
  [Range(2, 20)] public int gridResolution = 5;

  [Header("Camera Settings")]
  [Range(1, 4096)] public int cameraResolution = 256;
  public float nearClippingDistance = 0.01f;
  public LayerMask cullingMask = -1;
  public bool renderSkybox = false;

  [Header("Original Mesh")]
  public GameObject originalMesh;

  [HideInInspector] public Camera[] cameras;

  private GameObject impostorObject;
  private Material impostorMaterial;

  public bool HasImpostor => impostorObject != null;
  private float Radius => sphere_radius_ * transform.lossyScale.x;
  private const string OutputFolder = "Assets/yum_food/3ner/Impostor_Generated";

  void OnEnable() => Camera.onPreRender += UpdateMainCameraPos;
  void OnDisable() => Camera.onPreRender -= UpdateMainCameraPos;

  static void UpdateMainCameraPos(Camera cam)
  {
    if (cam.cameraType == CameraType.Game || cam.cameraType == CameraType.SceneView)
      Shader.SetGlobalVector("_ImpostorMainCameraPos", cam.transform.position);
  }

  void OnDrawGizmos()
  {
    Gizmos.color = Color.cyan;
    Gizmos.DrawWireSphere(transform.position, Radius);

    if (Application.isEditor && gridResolution > 0)
    {
      Gizmos.color = Color.yellow;
      for (int y = 0; y < gridResolution; y++)
      {
        for (int x = 0; x < gridResolution; x++)
        {
          Vector3 worldPos = transform.position + PlaneToHemiOctahedron(x, y) * (Radius + nearClippingDistance);
          Gizmos.DrawSphere(worldPos, Radius * 0.05f);
          Gizmos.DrawLine(worldPos, transform.position);
        }
      }
    }
  }

  Vector3 PlaneToHemiOctahedron(int gridX, int gridY)
  {
    float x = (gridX / (float)(gridResolution - 1)) * 2f - 1f;
    float z = (gridY / (float)(gridResolution - 1)) * 2f - 1f;

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

    // Octahedral decode
    float y = Mathf.Max(0f, 1f - Mathf.Abs(x_rot) - Mathf.Abs(z_rot));

    // Normalize
    Vector3 oct_pos = new Vector3(x_rot, y, z_rot).normalized;

    // Rotate back by -45° around y
    float rcp_sqrt2 = 0.70710678f;
    float x_unrot = (oct_pos.x - oct_pos.z) * rcp_sqrt2;
    float z_unrot = (oct_pos.x + oct_pos.z) * rcp_sqrt2;

    return new Vector3(x_unrot, oct_pos.y, z_unrot);
  }

  public void CreateCameras()
  {
    DestroyExistingCameras();

    GameObject parent = new GameObject("Cameras");
    parent.transform.SetParent(transform, false);

    cameras = new Camera[gridResolution * gridResolution];
    int idx = 0;

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

        GameObject camObj = new GameObject($"Camera_{x}_{y}");
        camObj.transform.SetParent(parent.transform, false);
        camObj.transform.position = worldPos;
        camObj.transform.LookAt(transform.position, transform.up);

        Camera cam = camObj.AddComponent<Camera>();
        cam.orthographic = true;
        cam.orthographicSize = sphere_radius_;
        cam.nearClipPlane = nearClippingDistance;
        cam.farClipPlane = sphere_radius_ * 2f + nearClippingDistance;
        cam.cullingMask = cullingMask;
        cam.clearFlags = renderSkybox ? CameraClearFlags.Skybox : CameraClearFlags.SolidColor;
        cam.backgroundColor = Color.clear;
        cam.enabled = false;

        cameras[idx++] = cam;
      }
    }
  }

  public void DestroyExistingCameras()
  {
    Transform camsTransform = transform.Find("Cameras");
    if (camsTransform != null) DestroyImmediate(camsTransform.gameObject);
    cameras = null;
  }

  public void BakeTexture()
  {
    SetRenderersEnabled(true);
    if (cameras == null || cameras.Length != gridResolution * gridResolution || cameras[0] == null) CreateCameras();

    // Load depth blit shader
    Shader depthBlitShader = Shader.Find("Hidden/yum_food/DepthBlit");
    if (depthBlitShader == null) { Debug.LogError("DepthBlit shader not found"); return; }
    Material depthBlitMat = new Material(depthBlitShader);

    // Render atlas
    int size = cameraResolution * gridResolution;
    Texture2D colorAtlas = new Texture2D(size, size, TextureFormat.RGBA32, false);
    Texture2D depthAtlas = new Texture2D(size, size, TextureFormat.RFloat, false);

    // Create RT with depth buffer that can be sampled as texture
    RenderTextureDescriptor desc = new RenderTextureDescriptor(cameraResolution, cameraResolution, RenderTextureFormat.ARGB32, 24);
    desc.sRGB = true;
    RenderTexture colorRT = RenderTexture.GetTemporary(desc);

    // Separate depth texture
    RenderTextureDescriptor depthDesc = new RenderTextureDescriptor(cameraResolution, cameraResolution, RenderTextureFormat.Depth, 24);
    RenderTexture depthOnlyRT = RenderTexture.GetTemporary(depthDesc);

    // Output for linearized depth
    RenderTexture linearDepthRT = RenderTexture.GetTemporary(cameraResolution, cameraResolution, 0, RenderTextureFormat.RFloat);

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

        // Render to color + depth buffers simultaneously
        cam.SetTargetBuffers(colorRT.colorBuffer, depthOnlyRT.depthBuffer);
        cam.Render();

        // Read color
        RenderTexture.active = colorRT;
        Texture2D colorTemp = new Texture2D(cameraResolution, cameraResolution);
        colorTemp.ReadPixels(new Rect(0, 0, cameraResolution, cameraResolution), 0, 0);
        colorTemp.Apply();
        colorAtlas.SetPixels(x * cameraResolution, y * cameraResolution, cameraResolution, cameraResolution, colorTemp.GetPixels());
        DestroyImmediate(colorTemp);

        // Blit depth buffer through linearization shader
        depthBlitMat.SetTexture("_DepthTex", depthOnlyRT);
        Graphics.Blit(null, linearDepthRT, depthBlitMat);

        // Read linearized depth
        RenderTexture.active = linearDepthRT;
        Texture2D depthTemp = new Texture2D(cameraResolution, cameraResolution, TextureFormat.RFloat, false);
        depthTemp.ReadPixels(new Rect(0, 0, cameraResolution, cameraResolution), 0, 0);
        depthTemp.Apply();
        depthAtlas.SetPixels(x * cameraResolution, y * cameraResolution, cameraResolution, cameraResolution, depthTemp.GetPixels());
        DestroyImmediate(depthTemp);
      }
    }
    colorAtlas.Apply();
    depthAtlas.Apply();
    RenderTexture.active = null;
    RenderTexture.ReleaseTemporary(colorRT);
    RenderTexture.ReleaseTemporary(depthOnlyRT);
    RenderTexture.ReleaseTemporary(linearDepthRT);
    DestroyImmediate(depthBlitMat);

    // Save
    if (!AssetDatabase.IsValidFolder(OutputFolder))
    {
      if (!AssetDatabase.IsValidFolder("Assets/yum_food/3ner"))
        AssetDatabase.CreateFolder("Assets/yum_food", "3ner");
      AssetDatabase.CreateFolder("Assets/yum_food/3ner", "Impostor_Generated");
    }

    string name = gameObject.name.Replace(" ", "_");
    string colorPath = Path.Combine(OutputFolder, $"{name}_atlas.png");
    string depthPath = Path.Combine(OutputFolder, $"{name}_depth.exr");

    File.WriteAllBytes(Path.Combine(Application.dataPath, "..", colorPath), colorAtlas.EncodeToPNG());
    File.WriteAllBytes(Path.Combine(Application.dataPath, "..", depthPath), depthAtlas.EncodeToEXR(Texture2D.EXRFlags.OutputAsFloat));
    DestroyImmediate(colorAtlas);
    DestroyImmediate(depthAtlas);

    AssetDatabase.Refresh();

    // Configure color atlas importer
    TextureImporter colorImporter = AssetImporter.GetAtPath(colorPath) as TextureImporter;
    if (colorImporter != null)
    {
      colorImporter.mipmapEnabled = true;
      colorImporter.alphaIsTransparency = true;
      colorImporter.wrapMode = TextureWrapMode.Clamp;
      colorImporter.filterMode = FilterMode.Trilinear;
      colorImporter.SaveAndReimport();
    }

    // Configure depth atlas importer
    TextureImporter depthImporter = AssetImporter.GetAtPath(depthPath) as TextureImporter;
    if (depthImporter != null)
    {
      depthImporter.mipmapEnabled = false;
      depthImporter.sRGBTexture = false;  // Linear data
      depthImporter.wrapMode = TextureWrapMode.Clamp;
      depthImporter.filterMode = FilterMode.Bilinear;
      depthImporter.textureCompression = TextureImporterCompression.Uncompressed;
      depthImporter.SaveAndReimport();
    }

    // Create impostor
    Texture2D colorTex = AssetDatabase.LoadAssetAtPath<Texture2D>(colorPath);
    Texture2D depthTex = AssetDatabase.LoadAssetAtPath<Texture2D>(depthPath);
    if (colorTex != null)
    {
      DestroyExistingImpostor();

      Shader shader = Shader.Find("yum_food/Gimmicks/Impostors");
      if (shader == null) { Debug.LogError("Shader not found"); return; }

      impostorMaterial = new Material(shader);
      impostorMaterial.SetTexture("_ImpostorAtlas", colorTex);
      impostorMaterial.SetTexture("_ImpostorDepthAtlas", depthTex);
      impostorMaterial.SetInt("_GridResolution", gridResolution);
      impostorMaterial.SetFloat("_SphereRadius", sphere_radius_);
      AssetDatabase.CreateAsset(impostorMaterial, Path.Combine(OutputFolder, $"{name}_mat.mat"));

      impostorObject = GameObject.CreatePrimitive(PrimitiveType.Quad);
      impostorObject.name = "Impostor";
      impostorObject.transform.SetParent(transform, false);
      impostorObject.transform.localScale = Vector3.one * sphere_radius_ * 2f;
      DestroyImmediate(impostorObject.GetComponent<Collider>());
      impostorObject.GetComponent<MeshRenderer>().sharedMaterial = impostorMaterial;

      SetRenderersEnabled(false);
    }
  }

  public void DestroyExistingImpostor()
  {
    if (impostorObject != null) DestroyImmediate(impostorObject);
    impostorObject = null;
    impostorMaterial = null;
  }

  public void ToggleRenderers()
  {
    if (originalMesh == null) return;
    bool showing = originalMesh.GetComponentInChildren<Renderer>()?.enabled ?? false;
    SetRenderersEnabled(!showing);
  }

  void SetRenderersEnabled(bool enabled)
  {
    if (originalMesh == null) return;
    foreach (Renderer r in originalMesh.GetComponentsInChildren<Renderer>(true))
      r.enabled = enabled;
    if (impostorObject != null) impostorObject.SetActive(!enabled);
  }
}

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

    GUILayout.Space(10);
    GUILayout.Label("Impostor Management", EditorStyles.boldLabel);

    EditorGUILayout.BeginHorizontal();
    if (GUILayout.Button("Create Impostor", GUILayout.Height(40))) s.BakeTexture();
    GUI.enabled = s.HasImpostor;
    if (GUILayout.Button("Destroy Impostor", GUILayout.Height(40)))
    {
      s.DestroyExistingImpostor();
      s.DestroyExistingCameras();
    }
    GUI.enabled = true;
    EditorGUILayout.EndHorizontal();

    if (s.originalMesh != null && GUILayout.Button("Toggle Visibility", GUILayout.Height(30)))
      s.ToggleRenderers();
  }
}