require 'yaml'

module BalazsVida
  module ActiveRecord
    module DataMigration
      def self.included(base)
        base.extend ClassMethods
      end
      
      module ClassMethods
          def create_from_yaml(klass_name)
            klass_name = klass_name.to_s if klass_name.kind_of?(Symbol)
            object_hash = {}
            
            say_with_time "Generate #{klass_name.camelize} from YAML" do
              klass_name = klass_name.camelize.constantize
              yaml_file = Dir["db/data/[0-9]*_#{self.name.to_s.tableize}.yml"].first
              data_hash = yaml_file ? YAML::load_file(yaml_file) : {} 
              data_hash.each do |key, data|
                obj = klass_name.new
                data.each do |attr_name, attr_value|
                  obj.send("#{attr_name}=", attr_value)
                end
                yield(key, obj) if block_given?
                obj.save! if obj.new_record?
                object_hash[key.to_sym] = obj
              end
            end
            
            object_hash
          end
          
          def generate_data(klass_name, range, &block)
            klass_name = klass_name.to_s if klass_name.kind_of?(Symbol)
            
            say_with_time "Generate mass of #{klass_name.camelize}" do
              klass_name = klass_name.camelize.constantize
        
              range.each do |index|
                obj = klass_name.new
                yield index, obj if block_given?
                obj.save! if obj.new_record?
              end
            end
          end
          
      end
    end
  end
end