require 'json'
require 'pathname'
require 'yaml'

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

HEADER = "This is free and unencumbered software released into the public domain."
BASE_URL = 'https://platform.openai.com'

Rust = Codify::Rust
OneOf = Rust::Types::Tuple0.new(:OneOf)
AnyOf = Rust::Types::Tuple0.new(:AnyOf)
AllOf = Rust::Types::Tuple0.new(:AllOf)

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

namespace :codegen do
  task debug: %w(openapi.yaml) 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 if definition2.primitive?
        p [definition2.name, definition2]
      end
    end
  end

  task components: %w(openapi.yaml) do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)

    File.open('src/components.rs', 'w') do |out|
      out.puts "// #{HEADER}"
      out.puts
      out.puts "//! OpenAI API components"
      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("components/#{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 if definition.primitive?
          out.puts
          definition.write(out)
        end
      end
    end
  end # :components

  task groups: %w(openapi.yaml) do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)

    File.open('src/groups.rs', 'w') do |out|
      out.puts "// #{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 "// #{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 :id, :ref, :type, :nullable, :default
    attr_reader :title, :description

    def initialize(id, spec)
      spec = spec.dup || {}
      @id = id ? id.to_sym : nil
      @ref = spec.delete('$ref')
      @type = spec['type'] ? spec.delete('type').to_sym : nil
      @nullable = spec.delete('nullable')
      @default = spec.delete('default')
      @title = spec.delete('title')
      @description = spec.delete('description')
      @spec = spec
    end

    def to_h
      @spec.dup.merge!({
        id: @id,
        '$ref': @ref,
        type: @type,
        nullable: @nullable,
        title: @title,
        description: @description,
      })
    end

    def nullable?() !!@nullable end
    def summary?() !self.summary.to_s.empty? end
    def summary() first_sentence(self.description) end

    def ref?() !!@ref end
    def recursive_ref?() !!@spec['$recursiveRef'] end
    def one_of?() !!@spec['oneOf'] end
    def any_of?() !!@spec['anyOf'] end
    def all_of?() !!@spec['allOf'] end
    def items?() !!@spec['items'] end
    def items() Schema.new("#{self.id}_Item", @spec['items']) end
    def properties?() !!@spec['properties'] end

    def one_of
      make_enums(@spec['oneOf'])
    end

    def any_of
      make_enums(@spec['anyOf'])
    end

    def all_of
      make_enums(@spec['allOf'])
    end

    def make_enums(inputs)
      if inputs && inputs.size == 1
        [Schema.new(self.id, inputs.first)]
      else
        (inputs || []).map.with_index { |x, i| Schema.new("#{self.id}_#{i + 1}", x) }
      end
    end

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

    def to_rust_variant(index = nil)
      result = case
        when self.ref? then self.ref.split('/').last
        when self.properties? || @type == :object
          'Object'
        when self.items? || @type == :array
          case self.items.type
            when :integer then 'ArrayOfIntegers'
            when :string then 'ArrayOfStrings'
            else 'Array'
          end
        when self.one_of? then 'OneOf'
        when self.any_of? then 'AnyOf'
        when self.all_of? then 'AllOf'
        else case @type
          when :boolean then 'Boolean'
          when :integer then 'Integer'
          when :number then 'Number'
          when :string then 'String'
          else raise NotImplementedError, self.inspect
        end
      end
      #index ? "#{result}#{index}" : result
      result
    end

    def to_rust(level = 0)
      type = case
        when self.ref?
          Rust::Types::Named.new(self.ref.split('/').last)
        when self.properties? || @type == :object
          # Some 414 top-level schemas are `type: "object"`.
          Rust::Struct.new(self.id, derives: %i(Clone Debug), cfg_derives: %i(serde)) do |definition|
            definition.comment = self.summary
            self.properties.each do |(property_id, property_type)|
              field_type = property_type.to_rust(level + 1)
              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
        when self.items? || self.type == :array
          # Only 4 top-level schemas are `type: "array"`.
          item_type = self.items.to_rust(level + 1) rescue Rust::Types::String
          Rust::Types::Vec.new(item_type)
        when self.one_of?
          # Some 23 top-level schemas are `oneOf`.
          schemas = self.one_of
          if schemas.size == 1
            schemas.first.to_rust
          elsif schemas.all? { |t| t.type == :string }
            # e.g. RealtimeResponseCreateParams_Conversation
            Rust::Newtype.new(self.id, Rust::Types::String, derives: %i(Clone Debug), cfg_derives: %i(serde)) do |definition|
              definition.comment = self.summary
            end
          elsif schemas.all? { |t| t.type == :object && t.properties.size == 2 && t.properties.any? { |k, _| k.to_sym == :type } }
            # e.g. CreateEvalCompletionsRunDataSource, CreateEvalResponsesRunDataSource
            Rust::Enum.new(self.id, derives: %i(Clone Debug), cfg_derives: %i(serde)) 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)
                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: %i(Clone Debug), cfg_derives: %i(serde)) do |definition|
              definition.comment = self.summary
              definition.variants << Rust::EnumVariant.new(:Null) if self.nullable?
              #out.puts "#[default]" if self.nullable? # TODO: && self.default == 'null'
              schemas.each_with_index do |t, i|
                next if t.recursive_ref? # skip $recursiveRef
                variant = Rust::EnumVariant.new(t.to_rust_variant(i), t.to_rust(level + 1))
                variant.comment = t.summary
                definition.variants << variant
              end
            end
          end
        when self.any_of?
          # Only 10 top-level schemas are `anyOf`.
          schemas = self.any_of
          if schemas.size == 1
            schemas.first.to_rust
          else
            Rust::Newtype.new(self.id, AnyOf, derives: %i(Clone Debug), cfg_derives: %i(serde)) do |definition|
              definition.comment = self.summary
            end
          end
        when self.all_of?
          # Only 9 top-level schemas are `allOf`.
          schemas = self.all_of
          if schemas.size == 1
            schemas.first.to_rust
          else
            Rust::Newtype.new(self.id, AllOf, derives: %i(Clone Debug), cfg_derives: %i(serde)) do |definition|
              definition.comment = self.summary
            end
          end
        else case @type
          # Only `ParallelToolCalls` has `type: "boolean"`.
          # TODO: spec['default'] if self.type == :boolean
          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 Rust::Types::String #raise NotImplementedError, self.inspect
        end
      end

      if level.zero? && !type.definition?
        type = Rust::Newtype.new(self.id, type, derives: %i(Clone Debug), cfg_derives: %i(serde)) do |definition|
          definition.comment = self.summary
        end
      end
      #type = @nullable ? Rust::Types::Option.new(type) : type # FIXME
      type
    end
  end # Schema

  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(YAML.load_file(path, aliases: true))
    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
      @spec['components']['schemas'].map { |k, v| OpenAPI::Schema.new(k, v) }.sort_by(&:id)
    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
