/*
 * 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;

namespace Firebase.AI.Internal {

// Include this to just shorthand it in this file
using JsonDict = Dictionary<string, object>;

// Options used by the extensions below to modify logic, mostly around throwing exceptions.
[Flags]
internal enum JsonParseOptions {
  None = 0,
  // Throw if the key is missing from the dictionary, useful for required fields.
  ThrowMissingKey = 1 << 0,
  // Throw if the key was found, but it failed to cast to the expected type.
  ThrowInvalidCast = 1 << 1,
  // Combination of the above.
  ThrowEverything = ThrowMissingKey | ThrowInvalidCast,
}

// Contains extension methods for commonly done logic.
internal static class FirebaseAIExtensions {

  // Tries to find the object with the key, and cast it to the type T
  public static bool TryParseValue<T>(this JsonDict jsonDict, string key,
      out T value, JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) {
    if (jsonDict.TryGetValue(key, out object obj)) {
      if (obj is T tObj) {
        value = tObj;
        return true;
      } else if (obj is long asLong &&
          (typeof(T) == typeof(int) ||
           typeof(T) == typeof(float) ||
           typeof(T) == typeof(double))) {
        // MiniJson puts all ints as longs, so special case it
        // Also, if the number didn't have a decimal point, we still want to
        // allow it to be a float or double.
        value = (T)Convert.ChangeType(asLong, typeof(T));
        return true;
      } else if (typeof(T) == typeof(float) && obj is double asDouble) {
        // Similarly, floats are stored as doubles.
        value = (T)(object)(float)asDouble;
        return true;
      } else if (options.HasFlag(JsonParseOptions.ThrowInvalidCast)) {
        throw new InvalidCastException(
            $"Invalid JSON format: '{key}' is not of type '{typeof(T).Name}', " +
            $"Actual type: {obj.GetType().FullName}.");
      }
    } else if (options.HasFlag(JsonParseOptions.ThrowMissingKey)) {
      throw new KeyNotFoundException(
          $"Invalid JSON format: Unable to locate expected key {key}");
    }
    value = default;
    return false;
  }

  // Casts the found object to type T, otherwise returns default (or throws)
  public static T ParseValue<T>(this JsonDict jsonDict, string key,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast,
      T defaultValue = default) {
    if (TryParseValue(jsonDict, key, out T value, options)) {
      return value;
    } else {
      return defaultValue;
    }
  }

  // Casts the found object to type T, otherwise returns null (or throws)
  public static T? ParseNullableValue<T>(this JsonDict jsonDict, string key,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) where T : struct {
    if (TryParseValue(jsonDict, key, out T value, options)) {
      return value;
    } else {
      return null;
    }
  }

  // Tries to convert the string found with the key to an enum, using the given function.
  public static bool TryParseEnum<T>(this JsonDict jsonDict, string key,
      Func<string, T> parseFunc, out T value,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) {
    if (jsonDict.TryParseValue(key, out string enumStr, options)) {
      value = parseFunc(enumStr);
      return true;
    } else {
      value = default;
      return false;
    }
  }

  // Casts the found string to an enum, otherwise returns default (or throws)
  public static T ParseEnum<T>(this JsonDict jsonDict, string key,
      Func<string, T> parseFunc,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast,
      T defaultValue = default) {
    if (TryParseEnum(jsonDict, key, parseFunc, out T value, options)) {
      return value;
    } else {
      return defaultValue;
    }
  }

  // Casts the found string to an enum, otherwise returns null (or throws)
  public static T? ParseNullableEnum<T>(this JsonDict jsonDict, string key,
      Func<string, T> parseFunc,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) where T : struct, Enum {
    if (TryParseEnum(jsonDict, key, parseFunc, out T value, options)) {
      return value;
    } else {
      return null;
    }
  }

  // Tries to convert the given Dictionary found at the key with the given function.
  public static bool TryParseObject<T>(this JsonDict jsonDict, string key,
      Func<JsonDict, T> parseFunc, out T value,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) {
    if (TryParseValue(jsonDict, key, out JsonDict innerDict, options)) {
      value = parseFunc(innerDict);
      return true;
    } else {
      value = default;
      return false;
    }
  }

  // Casts the found Dictionary to an object, otherwise returns default (or throws)
  public static T ParseObject<T>(this JsonDict jsonDict, string key,
      Func<JsonDict, T> parseFunc,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast,
      T defaultValue = default) {
    if (TryParseObject(jsonDict, key, parseFunc, out T value, options)) {
      return value;
    } else {
      return defaultValue;
    }
  }

  // Casts the found Dictionary to an object, otherwise returns null (or throws)
  public static T? ParseNullableObject<T>(this JsonDict jsonDict, string key,
      Func<JsonDict, T> parseFunc,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) where T : struct {
    if (TryParseObject(jsonDict, key, parseFunc, out T value, options)) {
      return value;
    } else {
      return null;
    }
  }

  // Tries to convert the found List of objects to a List of strings.
  public static bool TryParseStringList(this JsonDict jsonDict, string key,
      out List<string> values,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) {
    if (jsonDict.TryParseValue(key, out List<object> list, options)) {
      values = list.OfType<string>().ToList();
      return true;
    } else {
      values = null;
      return false;
    }
  }

  // Casts the found List to a string List, otherwise returns default (or throws)
  public static List<string> ParseStringList(this JsonDict jsonDict, string key,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) {
    TryParseStringList(jsonDict, key, out List<string> values, options);
    return values;
  }

  // Tries to convert the found List of Dictionaries, using the given function.
  public static bool TryParseObjectList<T>(this JsonDict jsonDict, string key,
      Func<JsonDict, T> parseFunc, out List<T> values,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) {
    if (jsonDict.TryParseValue(key, out List<object> list, options)) {
      values = list.ConvertJsonList(parseFunc);
      return true;
    } else {
      values = null;
      return false;
    }
  }

  // Casts the found List to an object List, otherwise returns default (or throws)
  public static List<T> ParseObjectList<T>(this JsonDict jsonDict, string key,
      Func<JsonDict, T> parseFunc,
      JsonParseOptions options = JsonParseOptions.ThrowInvalidCast) {
    TryParseObjectList(jsonDict, key, parseFunc, out List<T> values, options);
    return values;
  }

  // Converts the given List of Dictionaries into a list of T, using the converter.
  public static List<T> ConvertJsonList<T>(this List<object> list,
      Func<JsonDict, T> converter) {
    return list
        .Select(o => o as JsonDict)
        .Where(dict => dict != null)
        .Select(converter)
        .ToList();
  }

  public static ModelContent ConvertRole(this ModelContent content, string role) {
    if (content.Role == role) {
      return content;
    } else {
      return new ModelContent(role, content.Parts);
    }
  }

  public static ModelContent ConvertToUser(this ModelContent content) {
    return content.ConvertRole("user");
  }

  public static ModelContent ConvertToModel(this ModelContent content) {
    return content.ConvertRole("model");
  }

  public static ModelContent ConvertToSystem(this ModelContent content) {
    return content.ConvertRole("system");
  }
}
  
}
