require 'json'
require 'pathname'

require 'codify/rust' # https://rubygems.org/gems/codify.rb

FILE_HEADER = "This is free and unencumbered software released into the public domain."
BASE_URL = 'https://platform.openai.com'
OPENAPI_FILE = 'openapi.yaml'
DERIVES = %i(Clone Debug)
CFG_DERIVES = %i(serde)

Rust = Codify::Rust

task default: %w(codegen)
task codegen: %w(codegen:schemas codegen:groups)

require_relative '.rake/openapi'

include OpenAPI::Flattener
include OpenAPI::Transforms

namespace :codegen do
  task flatten: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)
    output = {}
    spec.schemas_raw.each do |schema_ref, schema|
      #next unless schema_ref == 'AssistantObject'
      flatten_definition(schema_ref, schema, output)
      # puts; puts JSON.pretty_generate(output)
      #break if schema_ref == 'AssistantObject'
    end
    [HoistNullables, RemoveRedundantCombines, DerefAliases, CombineAllOfs].each do |transform|
      transform.transform_schemas!(schemas)
    end
    output.each do |schema_ref, schema|
      #next unless schema[:oneOf] && schema[:oneOf].size < 3
      puts
      puts "// #{schema_ref}:"
      puts JSON.pretty_generate(schema)
    end
    p [spec.schemas_raw.keys.size, output.keys.size]
  end

  task debug: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)
    spec.schemas.each do |schema|
      #puts JSON.pretty_generate(schema.to_h)
      definition = schema.to_rust
      p [definition.name, definition]
      definition.each_subtype do |definition2|
        next unless definition2.definition?
        p [definition2.name, definition2]
      end
    end
  end

  task schemas: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)

    File.open('src/schemas.rs', 'w') do |out|
      out.puts "// #{FILE_HEADER}"
      out.puts
      out.puts "//! OpenAI API schemas"
      out.puts
      out.puts "#![allow(non_camel_case_types)]"
      out.puts
      out.puts "use crate::prelude::{String, Vec};"

      definitions = spec.schemas.map(&:to_rust)
      definitions.each do |definition|
        module_name = camel_to_snake(definition.name)
        module_path = Pathname("schemas/#{module_name}.rs")
        if (Pathname('src') / module_path).exist?
          out.puts
          out.puts "include!(\"#{module_path}\");"
        else
          out.puts
          definition.write(out)
        end

        definition.each_subtype do |definition|
          next unless definition.definition?
          out.puts
          definition.write(out)
        end
      end

      out.puts
      out.puts "include!(\"schemas/footer.rs\");"
    end
  end # :schemas

  task groups: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)

    File.open('src/groups.rs', 'w') do |out|
      out.puts "// #{FILE_HEADER}"
      out.puts
      out.puts "//! OpenAI API types organized by group"
      out.puts
      spec.groups.each do |group|
        module_name = group.snake_case_id
        out.puts "pub mod #{module_name};"
      end
    end

    spec.groups.each do |group|
      module_name = group.snake_case_id
      module_path = Pathname("src/groups/#{module_name}.rs")
      File.open(module_path, 'w') do |out|
        out.puts "// #{FILE_HEADER}"
        out.puts
        out.puts "//! **OpenAI API: #{group.title}**"

        next if group.description.empty?
        out.puts "//!"
        wrap_text(inline_to_reference_links(group.description), 80-4).each do |line|
          out.puts "//! #{line}"
        end

        link_refs = link_refs(group.description)
        next if link_refs.empty?
        out.puts "//!"
        link_refs.each do |link_ref|
          out.puts "//! #{link_ref}"
        end
      end
    end
  end # :groups
end # :codegen

module OpenAPI
  class Schema
    attr_reader :type, :id
    attr_reader :default
    attr_reader :title, :description

    ##
    # @param [Symbol, String] id
    # @param [Hash] spec
    def self.new(id, spec)
      raise ArgumentError, id.inspect unless id.is_a?(Symbol) || id.is_a?(String)
      raise ArgumentError, spec.inspect unless spec.is_a?(Hash)
      klass = case
        when spec[:'$recursiveRef']
          raise ArgumentError, "missing RemoveRedundantCombines transform: #{spec.inspect}"
        when spec[:allOf]
          raise ArgumentError, "missing CombineAllOfs transform: #{spec.inspect}"
        when spec[:'$ref'] then Ref
        when spec[:oneOf] || spec[:anyOf]
          choices = spec[:oneOf] || spec[:anyOf]
          return nil unless choices.is_a?(::Array) && choices.all? { it.is_a?(::Hash) } # FIXME
          OneOf
        when spec[:type] == 'object' || spec[:properties] then Object
        when spec[:type] == 'array' || spec[:items] then Array
        when spec[:type] == 'string' && spec[:enum] then StringEnum
        when spec[:type] && %i(null boolean integer number string).include?(spec[:type].to_sym) then Primitive
        else raise NotImplementedError, [id, spec].inspect
      end
      instance = klass.allocate
      instance.send(:initialize, id, spec)
      instance
    end

    def initialize(type, id, spec)
      spec.delete(:type)
      @type = type ? type.to_sym : nil
      @id = id ? id.to_sym : nil
      @default = spec.delete(:default)
      @nullable = spec.delete(:nullable)
      @title = spec.delete(:title)
      @description = spec.delete(:description)
      @deprecated = spec.delete(:deprecated)
      spec.delete(:example)
    end

    def defaultible?() self.nullable? end
    def nullable?() !!@nullable end
    def deprecated() !!@deprecated end

    def summary?() !self.summary.to_s.empty? end
    def summary() first_sentence(self.description) end
    def recursive_ref?() false end
  end # Schema

  class Ref < Schema
    attr_reader :ref

    def initialize(id, spec)
      spec = spec.dup || {}
      super(:ref, id, spec)
      raise ArgumentError, [:Ref, spec[:'$ref']].inspect unless spec[:'$ref']
      @ref = spec.delete(:'$ref').split('/').last
    end

    def defaultible?
      self.nullable?
    end

    def to_rust_variant
      self.ref
    end

    def to_rust(level = 0)
      type = Rust::Types::Named.new(self.ref)
      type = Rust::Types::Option.new(type) if self.nullable?
      if level.zero?
        Rust::TypeAlias.new(self.id, type) do |definition|
          definition.comment = self.summary
        end
      else
        type
      end
    end
  end # Ref

  ##
  # Schema of `type: "string"` with `enum: [...]` and (sometimes) `default: "..."`.
  class StringEnum < Schema
    def initialize(id, spec)
      spec = spec.dup || {}
      super(:string, id, spec)
      @enum = spec.delete(:enum)
      raise ArgumentError, [:StringEnum, @enum].inspect unless @enum.is_a?(::Array) && @enum.all? { |s| s.is_a?(String) }
    end

    def defaultible?
      false
    end

    def to_rust_variant
      :String
    end

    def to_rust(level = 0)
      if level.zero?
        derives = DERIVES + (self.defaultible? ? %i(Default) : [])
        Rust::Enum.new(self.id, derives: derives, cfg_derives: CFG_DERIVES, rename_all: :lowercase) do |definition|
          definition.comment = self.summary
          @enum.each do |variant_id|
            variant_name = snake_to_camel(Identifier.new(variant_id).to_rust)
            definition.variants << Rust::EnumVariant.new(variant_name, rename: variant_id)
          end
          definition.variants << Rust::EnumVariant.new(:Other, Rust::Types::String, other: true)
        end
      else
        # e.g. AssistantToolsFileSearchTypeOnly
        Rust::Types::String
      end
    end
  end # StringEnum

  class Primitive < Schema
    def initialize(id, spec)
      spec = spec.dup || {}
      super(spec.delete(:type), id, spec)
    end

    def defaultible?
      true
    end

    def to_rust_variant
      case @type
        when :null then :Null
        when :boolean then :Boolean
        when :integer then :Integer
        when :number then :Number
        when :string then :String
        else raise NotImplementedError, self.inspect
      end
    end

    def to_rust(level = 0)
      type = case @type
        # Only `ParallelToolCalls` has `type: "boolean"`.
        when :null then Rust::Types::Unit
        when :boolean then Rust::Types::Bool
        when :integer then Rust::Types::I64
        when :number then Rust::Types::F64
        # Some 12 schemas are `type: "string"`.
        when :string then Rust::Types::String
        else raise NotImplementedError, self.inspect
      end

      if level == 0
        # TODO: self.default
        Rust::TypeAlias.new(self.id, type) do |definition|
          definition.comment = self.summary
        end
      else
        type
      end
    end
  end # Primitive

  ##
  # Schema of `type: "object"` with `properties: {...}`.
  class Object < Schema
    def initialize(id, spec)
      spec = spec.dup || {}
      super(:object, id, spec)
      @properties = spec.delete(:properties) || {}
      @additional_properties = spec.delete(:additionalProperties)
      raise ArgumentError, [:Object, @properties].inspect unless @properties.is_a?(Hash)
      @required = spec.delete(:required)
      @required = @required ? @required.map(&:to_sym) : @required
      raise ArgumentError, [:Object, @required].inspect unless @required.nil? || @required.is_a?(::Array)
    end

    def defaultible?(level = 0)
      self.properties.all? do |property_id, property_type|
        is_option = !(@required&.include?(property_id.to_sym))
        is_option || property_type.defaultible? || property_type.to_rust(level + 1).defaultible?
      end
    end

    def properties
      result = @properties.inject({}) do |result, (k, v)|
        result[Identifier.new(k.to_sym)] = Schema.new("#{self.id}_#{snake_to_camel(k)}", v) # FIXME
        result
      end
      result.reject! { |_, v| v.nil? } # FIXME
      result
    end

    def to_rust_variant
      :Object
    end

    def to_rust(level = 0)
      # Some 414 top-level schemas are `type: "object"`.
      derives = DERIVES + (self.defaultible?(level) ? %i(Default) : [])
      Rust::Struct.new(self.id, derives: derives, cfg_derives: CFG_DERIVES) do |definition|
        definition.comment = self.summary
        self.properties.each do |(property_id, property_type)|
          is_option = !(@required&.include?(property_id.to_sym)) || property_type.nullable?
          field_type = property_type.to_rust(level + 1)
          field_type = Rust::Types::Option.new(field_type).flatten if is_option
          field = Rust::StructField.new(property_id.to_rust, field_type)
          field.comment = property_type.summary
          field.rename = property_id.id if property_id.needs_rename?
          definition.fields << field
        end
      end
    end
  end # Object

  ##
  # Schema of `type: "array"` with `items: [...]`.
  class Array < Schema
    def initialize(id, spec)
      spec = spec.dup || {}
      super(:array, id, spec)
      @items = spec.delete(:items)
      raise ArgumentError, [:Array, @items].inspect unless @items.is_a?(Hash)
    end

    def defaultible?
      self.items.defaultible?
    end

    def items?() !!@items end
    def items() Schema.new(self.id, @items) end

    def to_rust_variant
      case self.items.type
        when :integer then :ArrayOfIntegers
        when :string then :ArrayOfStrings
        else :Array
      end
    end

    def to_rust(level = 0)
      # Only 4 top-level schemas are `type: "array"`.
      item_type = self.items.to_rust(level + 1) #rescue Rust::Types::String # FIXME
      type = Rust::Types::Vec.new(item_type)
      if level.zero?
        Rust::TypeAlias.new(self.id, type) do |definition|
          definition.comment = self.summary
        end
      else
        type
      end
    end
  end # Array

  class Combine < Schema
    attr_reader :choices
  end

  ##
  # Schema of `oneOf: [...]` or `anyOf: [...]`.
  class OneOf < Combine
    TYPE = Rust::Types::Tuple0.new(:OneOf)

    def initialize(id, spec)
      spec = spec.dup || {}
      super(:one_of, id, spec)
      @choices = spec.delete(:oneOf) || spec.delete(:anyOf)
      raise ArgumentError, [:OneOf, @choices].inspect unless @choices.is_a?(::Array) && @choices.all? { it.is_a?(::Hash) }
      raise ArgumentError, [:OneOf, @choices].inspect unless @choices.size > 1
      @choices.map!.with_index { |x, i| Schema.new("#{self.id}_#{i + 1}", x) }
    end

    def defaultible?
      self.nullable?
    end

    def to_rust_variant
      :OneOf
    end

    def to_rust(level = 0)
      # Some 23 top-level schemas are `oneOf`. Only 10 top-level schemas are `anyOf`.
      schemas = self.choices
      if schemas.all? { |t| t.type.to_sym == :string }
        # e.g. RealtimeResponseCreateParams_Conversation
        Rust::TypeAlias.new(self.id, Rust::Types::String) do |definition|
          definition.comment = self.summary
        end
        # TODO: Rust::Types::String
      elsif schemas.all? { |t| t.type.to_sym == :string || t.type == :null }
        # e.g. ApproximateLocation_City
        Rust::TypeAlias.new(self.id, Rust::Types::Option.new(Rust::Types::String)) do |definition|
          definition.comment = self.summary
        end
        # TODO: Rust::Types::Option.new(Rust::Types::String)
      elsif schemas.size == 2 && schemas.any? { |t| t && t.type.to_sym == :null }
        type = schemas.find { |t| t.type.to_sym != :null }
        Rust::TypeAlias.new(self.id, Rust::Types::Option.new(type.to_rust)) do |definition|
          definition.comment = self.summary
        end
        # TODO: Rust::Types::Option.new(type.to_rust)
      # elsif schemas.all? { |t| t.type.to_sym == :object && t.properties.size == 2 && t.properties.any? { |k, _| k.to_sym == :type } }
      #   # e.g. CreateEvalCompletionsRunDataSource, CreateEvalResponsesRunDataSource
      #   Rust::Enum.new(self.id, derives: DERIVES, cfg_derives: CFG_DERIVES) do |definition|
      #     definition.comment = self.summary
      #     schemas.each do |t|
      #       property1, property2 = t.properties.values
      #       variant_name = snake_to_camel(property1.to_h[:enum].first) # FIXME!
      #       variant_type = property2.to_rust(level + 1)
      #       variant = Rust::EnumVariant.new(variant_name, variant_type)
      #       variant.comment = t.summary
      #       definition.variants << variant
      #     end
      #   end
      else
        Rust::Enum.new(self.id, derives: DERIVES, cfg_derives: CFG_DERIVES) do |definition|
          definition.comment = self.summary
          definition.variants << Rust::EnumVariant.new(:Null, default: true) if self.nullable?
          schemas.each do |t|
            variant = Rust::EnumVariant.new(t.to_rust_variant, t.to_rust(level + 1))
            variant.comment = t.summary
            definition.variants << variant
          end
        end
      end
    end
  end # OneOf

  class Identifier
    include Comparable

    attr_reader :id

    def initialize(id) @id = id end
    def <=>(other) self.id <=> other.id end
    def ==(other) self.id == other.id end
    def hash() @id.hash end
    def needs_rename?() @id.to_s.match?(/[\[\]\.]/) end
    def to_s() @id.to_s end
    def to_sym() @id.to_sym end
    def to_rust() @id.to_s.gsub('[]', '').gsub(/[-\/\.]/, '_') end
  end # Identifier
end

module OpenAI
  class Spec
    def self.parse(path)
      self.new(OpenAPI.load_file(path))
    end

    def initialize(spec) @spec = spec end

    def groups
      @spec[:'x-oaiMeta'][:groups].map { |g| Group.new(g) }.sort_by(&:snake_case_id)
    end

    def schemas
      schemas = flatten_openapi_schemas(self.schemas_raw)
      [HoistNullables, RemoveRedundantCombines, DerefAliases, CombineAllOfs].each do |transform|
        transform.transform_schemas!(schemas)
      end
      schemas.map { |k, v| OpenAPI::Schema.new(k, v) }.sort_by(&:id)
    end

    def schemas_raw
      @spec[:components][:schemas]
    end
  end # Spec

  class Group
    attr_reader :spec

    def initialize(spec) @spec = spec end
    def snake_case_id() @spec[:id].gsub('-', '_') end
    def title() @spec[:title] end
    def description() @spec[:description] end
  end
end # OpenAI

def wrap_text(text, max_width = 80)
  text.to_s.gsub(/(.{1,#{max_width}})(\s+|$)/, "\\1\n").split("\n")
end

def first_sentence(text)
  text.to_s.gsub(/\s*\n+/, ' ').match(/^.*?[.!?](?:\s|$)/)&.[](0)&.strip || text.to_s.strip
end

def link_refs(markdown, base_url = BASE_URL)
  base_url = base_url.to_s.chomp('/')
  markdown.to_s.scan(/\[(.*?)\]\((\/[^)]*)\)/)
    .map { |_, path| "[#{path}]: #{base_url}#{path}" }
    .uniq
end

def inline_to_reference_links(markdown)
  markdown.to_s.gsub(/\[(.*?)\]\(((?!https:).*?)\)/, '[\1][\2]')
end

def camel_to_snake(name)
  name.to_s.gsub(/([A-Z])/, '_\1').downcase.gsub(/^_/, '')
end

def snake_to_camel(name)
  name.to_s.gsub(/[_.]/, ' ').split.map(&:capitalize).join
end
