/*
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

using System;
using System.Collections.Generic;
using System.Linq;
using Firebase.AI.Internal;
using Google.MiniJSON;
using UnityEngine;

namespace Firebase.AI
{
  /// <summary>
  /// An image generated by Imagen.
  /// </summary>
  public interface IImagenImage
  {
    /// <summary>
    /// The IANA standard MIME type of the image file; either `"image/png"` or `"image/jpeg"`.
    ///
    /// > Note: To request a different format, set `ImageFormat` in
    ///   your `ImagenGenerationConfig`.
    /// </summary>
    public string MimeType { get; }
  }

  /// <summary>
  /// An image generated by Imagen, represented as inline data.
  /// </summary>
  public readonly struct ImagenInlineImage : IImagenImage
  {
    /// <summary>
    /// The IANA standard MIME type of the image file; either `"image/png"` or `"image/jpeg"`.
    ///
    /// > Note: To request a different format, set `ImageFormat` in
    ///   your `ImagenGenerationConfig`.
    /// </summary>
    public string MimeType { get; }
    /// <summary>
    /// The image data in PNG or JPEG format.
    /// </summary>
    public byte[] Data { get; }

    /// <summary>
    /// Convert the image data into a `UnityEngine.Texture2D`.
    /// </summary>
    /// <returns></returns>
    public UnityEngine.Texture2D AsTexture2D()
    {
      var texture = new Texture2D(1, 1);
      texture.LoadImage(Data);
      return texture;
    }

    private ImagenInlineImage(string mimeType, byte[] data)
    {
      MimeType = mimeType;
      Data = data;
    }

    /// <summary>
    /// Intended for internal use only.
    /// This method is used for deserializing JSON responses and should not be called directly.
    /// </summary>
    internal static IImagenImage FromJson(Dictionary<string, object> jsonDict)
    {
      return new ImagenInlineImage(
        jsonDict.ParseValue<string>("mimeType", JsonParseOptions.ThrowEverything),
        Convert.FromBase64String(jsonDict.ParseValue<string>("bytesBase64Encoded", JsonParseOptions.ThrowEverything)));
    }
  }

  /// <summary>
  /// A response from a request to generate images with Imagen.
  ///
  /// This type is returned from:
  ///   - `ImagenModel.GenerateImagesAsync(prompt)` where `T` is `ImagenInlineImage`
  /// </summary>
  public readonly struct ImagenGenerationResponse<T> where T : IImagenImage
  {
    /// <summary>
    /// The images generated by Imagen; see `ImagenInlineImage`.
    ///
    /// > Important: The number of images generated may be fewer than the number requested if one or
    ///   more were filtered out; see `FilteredReason`.
    /// </summary>
    public IReadOnlyList<T> Images { get; }
    /// The reason, if any, that generated images were filtered out.
    ///
    /// This property will only be populated if fewer images were generated than were requested,
    /// otherwise it will be `null`. Images may be filtered out due to the `SafetyFilterLevel`,
    /// the `PersonFilterLevel`, or filtering included in the model. The filter levels may be
    /// adjusted in your `ImagenSafetySettings`. See the [Responsible AI and usage guidelines for
    /// Imagen](https://cloud.google.com/vertex-ai/generative-ai/docs/image/responsible-ai-imagen)
    /// for more details.
    public string FilteredReason { get; }

    private ImagenGenerationResponse(List<T> images, string filteredReason)
    {
      Images = images ?? new List<T>();
      FilteredReason = filteredReason;
    }

    /// <summary>
    /// Intended for internal use only.
    /// This method is used for deserializing JSON responses and should not be called directly.
    /// </summary>
    internal static ImagenGenerationResponse<T> FromJson(string jsonString)
    {
      return FromJson(Json.Deserialize(jsonString) as Dictionary<string, object>);
    }

    /// <summary>
    /// Intended for internal use only.
    /// This method is used for deserializing JSON responses and should not be called directly.
    /// </summary>
    internal static ImagenGenerationResponse<T> FromJson(Dictionary<string, object> jsonDict)
    {
      if (!jsonDict.ContainsKey("predictions") || jsonDict["predictions"] is not List<object>)
      {
        return new ImagenGenerationResponse<T>(null,
            "Model response missing predictions");
      }

      var predictions = (jsonDict["predictions"] as List<object>)
          .OfType<Dictionary<string, object>>();

      List<T> images = new();
      string filteredReason = null;

      foreach (var pred in predictions)
      {
        if (pred.ContainsKey("bytesBase64Encoded") && typeof(T) == typeof(ImagenInlineImage))
        {
          images.Add((T)ImagenInlineImage.FromJson(pred));
        }
        else if (pred.ContainsKey("raiFilteredReason"))
        {
          filteredReason = pred.ParseValue<string>("raiFilteredReason", JsonParseOptions.ThrowEverything);
        }
      }

      return new ImagenGenerationResponse<T>(images, filteredReason);
    }
  }
}
