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

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

namespace :codegen do
  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 "use crate::prelude::{String, Vec};"

      spec.schemas.each do |schema|
        out.puts

        module_name = camel_to_snake(schema.id.to_s)
        module_path = Pathname("components/#{module_name}.rs")

        if (Pathname('src') / module_path).exist?
          out.puts "include!(\"#{module_path}\");"
          next
        end

        definition = case schema.type
          when :boolean, :string, :array
            # Only `ParallelToolCalls` has a type of 'boolean'
            # Only 4 schemas have a type of 'array'
            # Some 12 schemas have a type of 'string'
            case schema.type
              when :boolean then Rust::Newtype.new(schema.id, :bool)
              when :string then Rust::Newtype.new(schema.id, :String)
              when :array then Rust::Newtype.new(schema.id, "Vec<#{schema.items.to_rust}>")
            end
            # TODO: schema['default'] if schema.type == :boolean
          when :object
            definition = Rust::Struct.new(schema.id, derives: %i(Clone Debug))
            definition.comment = schema.summary
            schema.properties.each do |(k, v)|
              field = Rust::Field.new(k, v)
              field.comment = v.summary
              field.rename = k.id if k.needs_rename?
              definition.fields << field
            end
            definition
          when nil then case
            # Some 49 schemas have a type of 'nil` with `oneOf`, `anyOf`, `allOf`, or `properties`
            when schema.one_of?
              #puts JSON.pretty_generate(schema.to_h)
              definition = Rust::Enum.new(schema.id, derives: %i(Clone Debug))
              # out.puts "#[default]" if schema.nullable? # TODO: && schema.default == 'null'
              definition.variants << Rust::Variant.new(:Null) if schema.nullable?
              schema.one_of.each do |t|
                variant = Rust::Variant.new(t.to_rust_variant, t.to_rust)
                variant.comment = t.summary
                definition.variants << variant
              end
              definition
            when schema.any_of?
              Rust::Newtype.new(schema.id, '(/*AnyOf*/)', derives: %i(Clone Debug)) # TODO
            when schema.all_of?
              Rust::Newtype.new(schema.id, '(/*AllOf*/)', derives: %i(Clone Debug)) # TODO
            when schema.properties?
              Rust::Struct.new(schema.id, derives: %i(Clone Debug)) # TODO
            else raise NotImplementedError # shouldn't reach here
          end
        end

        definition.comment = schema.summary
        definition.derives << :Default if false # TODO
        definition.write(out)
      end
    end
  end

  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 Type
    # def self.new(spec)
    #   Type.new(nil, spec)
    # end

    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 one_of?() !!@spec['oneOf'] end
    def any_of?() !!@spec['anyOf'] end
    def all_of?() !!@spec['allOf'] end
    def items?() !!@spec['items'] end
    def items() Type.new(nil, @spec['items']) end
    def properties?() !!@spec['properties'] end

    def one_of
      (@spec['oneOf'] || []).map { |x| Type.new(nil, x) }
    end

    def any_of
      (@spec['anyOf'] || []).map { |x| Type.new(nil, x) }
    end

    def all_of
      (@spec['allOf'] || []).map { |x| Type.new(nil, x) }
    end

    def properties
      (@spec['properties'] || {}).inject({}) do |result, (k, v)|
        result[Identifier.new(k)] = Type.new(nil, v)
        result
      end
    end

    def to_rust_variant
      case
        when self.ref? then self.ref.split('/').last
        when self.one_of? then 'OneOf'
        when self.any_of? then 'AnyOf'
        when self.all_of? then 'AllOf'
        when self.items? || @type == :array then 'Array'
        when self.properties? || @type == :object then 'Object'
        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
    end

    def to_rust_basetype
      case
        when self.ref? then self.ref.split('/').last
        when self.one_of? then '(/*OneOf*/)' # TODO
        when self.any_of? then '(/*AnyOf*/)' # TODO
        when self.all_of? then '(/*AllOf*/)' # TODO
        when self.items? || @type == :array then "Vec<#{self.items.to_rust}>"
        when self.properties? || @type == :object then '(/*Object*/)' # TODO
        else case @type
          when :boolean then :bool
          when :integer then :i64
          when :number then :f64
          when :string then :String
          else raise NotImplementedError, self.inspect
        end
      end
    end

    def to_rust
      basetype = self.to_rust_basetype
      @nullable ? "Option<#{basetype}>" : basetype
    end
  end

  class Identifier
    include Comparable

    attr_reader :id

    def initialize(id) @id = id end
    def <=>(other) self.id <=> other.id 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
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::Type.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

module Rust
  class Definition
    attr_reader :name
    attr_reader :derives
    attr_accessor :comment

    def initialize(name, derives: nil, comment: nil)
      @name = name
      @derives = (derives || []).to_a.dup.uniq.sort
      @comment = comment ? comment.to_s.strip : nil
    end

    def comment?() self.comment && !self.comment.empty? end

    def write(out)
      out.puts wrap_text(self.comment, 80-4).map { |s| s.prepend("/// ") }.join("\n") if self.comment?
      out.puts "#[derive(#{@derives.sort.join(", ")})]" unless @derives.empty?
      out.puts "#[cfg_attr(feature = \"serde\", derive(serde::Serialize, serde::Deserialize))]"
    end
  end # Definition

  class Newtype < Definition
    def initialize(name, type, derives: %i(Clone Debug), comment: nil)
      super(name, derives:, comment:)
      @type = type.to_s
    end

    def write(out)
      super(out)
      out.puts "pub struct #{@name}(pub #{@type});"
    end
  end # Struct

  class Struct < Definition
    attr_reader :fields

    def initialize(name, fields: nil, derives: nil, comment: nil)
      super(name, derives:, comment:)
      @fields = (fields || []).to_a.dup
    end

    def write(out)
      super(out)
      if self.fields.empty?
        out.puts "pub struct #{@name};"
      else
        out.puts "pub struct #{@name} {"
        @fields.each_with_index do |field, i|
          out.puts if i > 0
          out.puts wrap_text(field.comment, 80-8).map { |s| s.prepend("    /// ") }.join("\n") if field.comment?
          out.puts "    #[cfg_attr(feature = \"serde\", serde(rename = \"#{k.id}\"))]" if false
          field.write(out)
        end
        out.puts "}"
      end
    end
  end # Struct

  class Enum < Definition
    attr_reader :variants

    def initialize(name, variants: nil, derives: nil, comment: nil)
      super(name, derives:, comment:)
      @variants = (variants || []).to_a.dup
    end

    def write(out)
      super(out)
      if self.variants.empty?
        out.puts "pub struct #{@name};"
      else
        out.puts "pub enum #{@name} {"
        @variants.each_with_index do |variant, i|
          out.puts if i > 0
          out.puts wrap_text(variant.comment, 80-8).map { |s| s.prepend("    /// ") }.join("\n") if variant.comment?
          variant.write(out)
        end
        out.puts "}"
      end
    end
  end # Enum

  class Field
    attr_reader :name, :type, :summary
    attr_accessor :comment, :rename

    def initialize(name, type)
      @name = name.to_s.gsub('[]', '').gsub('.', '_')
      @type = type.to_rust # FIXME
    end

    def comment?() self.comment && !self.comment.empty? end

    def write(out)
      out.puts "    #[cfg_attr(feature = \"serde\", serde(rename = \"#{self.rename}\"))]" if self.rename
      out.puts "    pub r##{@name}: #{@type},"
    end
  end # Field

  class Variant
    attr_reader :name, :type, :summary
    attr_accessor :comment

    def initialize(name, type = nil)
      @name = name.to_sym
      @type = type ? type.to_s : nil
    end

    def comment?() self.comment && !self.comment.empty? end

    def write(out)
      if !@type
        out.puts "    #{@name},"
      else
        out.puts "    #{@name}(#{@type}),"
      end
    end
  end # Variant
end # Rust

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

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

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

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

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