module Mirah  
  module JVM
    module Compiler
      class JVMBytecode < Base
        java_import java.lang.System
        java_import java.io.PrintStream
        include Mirah::JVM::MethodLookup
        Types = Mirah::JVM::Types
        
        class << self
          attr_accessor :verbose
          
          def log(message)
            puts "* [#{name}] #{message}" if JVMBytecode.verbose
          end
          
          def classname_from_filename(filename)
            basename = File.basename(filename).sub(/\.(duby|mirah)$/, '')
            basename.split(/[_-]/).map{|x| x[0...1].upcase + x[1..-1]}.join
          end
        end
        
        module JVMLogger
          def log(message); JVMBytecode.log(message); end
        end
        
        class ImplicitSelf
          attr_reader :inferred_type
          
          def initialize(type)
            @inferred_type = type
          end
          
          def compile(compiler, expression)
            compiler.compile_self if expression
          end
        end
        
        def initialize
          super
          @jump_scope = []
        end
        
        def file_builder(filename)
          builder = BiteScript::FileBuilder.new(filename)
          AST.type_factory.define_types(builder)
          builder
        end
        
        def output_type
          "classes"
        end
        
        def push_jump_scope(node)
          raise "Not a node" unless Mirah::AST::Node === node
          begin
            @jump_scope << node
            yield
          ensure
            @jump_scope.pop
          end
        end
        
        def find_ensures(before)
          found = []
          @jump_scope.reverse_each do |scope|
            if Mirah::AST::Ensure === scope
              found << scope
            end
            break if before === scope
          end
          found
        end
        
        def begin_main
          # declare argv variable
          @method.local('argv', AST.type(nil, 'string', true))
        end
        
        def finish_main
          @method.returnvoid
        end
        
        def prepare_binding(scope)
          if scope.has_binding?
            type = scope.binding_type
            @binding = @bindings[type]
            @method.new type
            @method.dup
            @method.invokespecial type, "<init>", [@method.void]
            if scope.respond_to? :arguments
              scope.arguments.args.each do |param|
                name = param.name
                param_type = param.inferred_type
                if scope.static_scope.captured?(param.name)
                  @method.dup
                  type.load(@method, @method.local(name, param_type))
                  @method.putfield(type, name, param_type)
                end
              end
            end
            type.store(@method, @method.local('$binding', type))
          end
          begin
            yield
          ensure
            if scope.has_binding?
              @binding.stop
              @binding = nil
            end
          end
        end
        
        def define_method(node)
          push_jump_scope(node) do
            base_define_method(node, true) do |method, arg_types|
              return if @class.interface?
              
              log "Starting new #{node.static? ? 'static ' : ''}method #{node.name}(#{arg_types})"
              args = node.arguments.args
              method_body(method, args, node, node.signature[:return])
              log "Method #{node.name}(#{arg_types}) complete!"
            end
          end
        end
        
        def define_optarg_chain(name, arg, return_type,
          args_for_opt, arg_types_for_opt)
          # declare all args so they get their values
          @method.aload(0) unless @static
          args_for_opt.each do |req_arg|
            req_arg.inferred_type.load(@method, @method.local(req_arg.name, req_arg.inferred_type))
          end
          arg.value.compile(self, true)
          
          # invoke the next one in the chain
          if @static
            @method.invokestatic(@class, name.to_s, [return_type] + arg_types_for_opt + [arg.inferred_type])
          else
            @method.invokevirtual(@class, name.to_s, [return_type] + arg_types_for_opt + [arg.inferred_type])
          end
          
          return_type.return(@method)
        end
        
        def constructor(node)
          push_jump_scope(node) do
            super(node, true) do |method, args|
              method_body(method, args, node, Types::Void) do
                method.aload 0
                if node.delegate_args
                  if node.calls_super
                    delegate_class = @type.superclass
                  else
                    delegate_class = @type
                  end
                  delegate_types = node.delegate_args.map do |arg|
                    arg.inferred_type
                  end
                  constructor = delegate_class.constructor(*delegate_types)
                  node.delegate_args.each do |arg|
                    arg.compile(self, true)
                  end
                  method.invokespecial(
                  delegate_class, "<init>",
                  [@method.void, *constructor.argument_types])
                else
                  method.invokespecial @class.superclass, "<init>", [@method.void]
                end
              end
            end
          end
        end
        
        def method_body(method, args, node, return_type)
          body = node.body
          with(:method => method,
          :declared_locals => {}) do
            
            method.start
            
            scope = node.static_scope
            
            # declare all args so they get their values
            if args
              args.each {|arg| declare_local(scope, arg.name, arg.inferred_type)}
            end
            declare_locals(scope)
            
            yield if block_given?
            
            prepare_binding(node) do
              expression = return_type != Types::Void
              body.compile(self, expression) if body
            end
            
            return_type.return(@method)
            
            @method.stop
          end
        end
        
        def define_closure(class_def, expression)
          compiler = ClosureCompiler.new(@file, @type, self)
          compiler.define_class(class_def, expression)
        end
        
        def branch(iff, expression)
          elselabel = @method.label
          donelabel = @method.label
          
          # this is ugly...need a better way to abstract the idea of compiling a
          # conditional branch while still fitting into JVM opcodes
          predicate = iff.condition.predicate
          if iff.body || expression
            jump_if_not(predicate, elselabel)
            
            if iff.body
              iff.body.compile(self, expression)
            elsif expression
              iff.inferred_type.init_value(@method)
            end
            
            @method.goto(donelabel)
          else
            jump_if(predicate, donelabel)
          end
          
          elselabel.set!
          
          if iff.else
            iff.else.compile(self, expression)
          elsif expression
            iff.inferred_type.init_value(@method)
          end
          
          donelabel.set!
        end
        
        def loop(loop, expression)
          push_jump_scope(loop) do
            with(:break_label => @method.label,
            :redo_label => @method.label,
            :next_label => @method.label) do
              predicate = loop.condition.predicate
              
              loop.init.compile(self, false) if loop.init?
              
              pre_label = @redo_label
              
              if loop.check_first
                @next_label.set! unless loop.post?
                if loop.negative
                  # if condition, exit
                  jump_if(predicate, @break_label)
                else
                  # if not condition, exit
                  jump_if_not(predicate, @break_label)
                end
              end
              
              if loop.pre?
                pre_label = method.label
                pre_label.set!
                loop.pre.compile(self, false)
              end
              
              
              @redo_label.set!
              loop.body.compile(self, false) if loop.body
              
              if loop.check_first && !loop.post?
                @method.goto(@next_label)
              else
                @next_label.set!
                loop.post.compile(self, false) if loop.post?
                if loop.negative
                  # if not condition, continue
                  jump_if_not(predicate, pre_label)
                else
                  # if condition, continue
                  jump_if(predicate, pre_label)
                end
              end
              
              @break_label.set!
              
              # loops always evaluate to null
              @method.aconst_null if expression
            end
          end
        end
        
        def break(node)
          error("break outside of loop", node) unless @break_label
          handle_ensures(find_ensures(Mirah::AST::Loop))
          @method.goto(@break_label)
        end
        
        def next(node)
          error("next outside of loop", node) unless @next_label
          handle_ensures(find_ensures(Mirah::AST::Loop))
          @method.goto(@next_label)
        end
        
        def redo(node)
          error("redo outside of loop", node) unless @redo_label
          handle_ensures(find_ensures(Mirah::AST::Loop))
          @method.goto(@redo_label)
        end
        
        def jump_if(predicate, target)
          unless predicate.inferred_type == Types::Boolean
            raise "Expected boolean, found #{predicate.inferred_type}"
          end
          if Mirah::AST::Call === predicate
            method = extract_method(predicate)
            if method.respond_to? :jump_if
              method.jump_if(self, predicate, target)
              return
            end
          end
          predicate.compile(self, true)
          @method.ifne(target)
        end
        
        def jump_if_not(predicate, target)
          unless predicate.inferred_type == Types::Boolean
            raise "Expected boolean, found #{predicate.inferred_type}"
          end
          if Mirah::AST::Call === predicate
            method = extract_method(predicate)
            if method.respond_to? :jump_if_not
              method.jump_if_not(self, predicate, target)
              return
            end
          end
          predicate.compile(self, true)
          @method.ifeq(target)
        end
        
        def extract_method(call)
          target = call.target.inferred_type!
          params = call.parameters.map do |param|
            param.inferred_type!
          end
          target.get_method(call.name, params)
        end
        
        def call(call, expression)
          return cast(call, expression) if call.cast?
          method = extract_method(call)
          if method
            method.call(self, call, expression)
          else
            raise "Missing method #{target}.#{call.name}(#{params.join ', '})"
          end
        end
        
        def self_call(fcall, expression)
          return cast(fcall, expression) if fcall.cast?
          type = fcall.scope.static_scope.self_type
          type = type.meta if (@static && type == @type)
          fcall.target = ImplicitSelf.new(type)
          
          params = fcall.parameters.map do |param|
            param.inferred_type
          end
          method = type.get_method(fcall.name, params)
          unless method
            target = static ? @class.name : 'self'
            
            raise NameError, "No method %s.%s(%s)" %
            [target, fcall.name, params.join(', ')]
          end
          method.call(self, fcall, expression)
        end
        
        def super_call(sup, expression)
          type = @type.superclass
          sup.target = ImplicitSelf.new(type)
          
          params = sup.parameters.map do |param|
            param.inferred_type
          end
          method = type.get_method(sup.name, params)
          unless method
            
            raise NameError, "No method %s.%s(%s)" %
            [type, sup.name, params.join(', ')]
          end
          method.call_special(self, sup, expression)
        end
        
        def cast(fcall, expression)
          # casting operation, not a call
          castee = fcall.parameters[0]
          
          # TODO move errors to inference phase
          source_type_name = castee.inferred_type.name
          target_type_name = fcall.inferred_type.name
          if castee.inferred_type.primitive?
            if fcall.inferred_type.primitive?
              if source_type_name == 'boolean' && target_type_name != "boolean"
                raise TypeError.new "not a boolean type: #{castee.inferred_type}"
              end
              # ok
              primitive = true
            else
              raise TypeError.new "Cannot cast #{castee.inferred_type} to #{fcall.inferred_type}: not a reference type."
            end
          elsif fcall.inferred_type.primitive?
            raise TypeError.new "not a primitive type: #{castee.inferred_type}"
          else
            # ok
            primitive = false
          end
          
          castee.compile(self, expression)
          if expression
            if primitive
              source_type_name = 'int' if %w[byte short char].include? source_type_name
              if (source_type_name != 'int') && (%w[byte short char].include? target_type_name)
                target_type_name = 'int'
              end
              
              if source_type_name != target_type_name
                if RUBY_VERSION == "1.9"
                  @method.send "#{source_type_name[0]}2#{target_type_name[0]}"
                else
                  @method.send "#{source_type_name[0].chr}2#{target_type_name[0].chr}"
                end
              end
            else
              if (source_type_name != target_type_name ||
                castee.inferred_type.array? != fcall.inferred_type.array?)
                @method.checkcast fcall.inferred_type
              end
            end
          end
        end
        
        def body(body, expression)
          # last element is an expression only if the body is an expression
          super(body, expression) do |last|
            compile(last, expression)
          end
        end
        
        def local(scope, name, type)
          type.load(@method, @method.local(scoped_local_name(name, scope), type))
        end
        
        def local_assign(scope, name, type, expression, value)
          declare_local(scope, name, type)
          
          value.compile(self, true)
          
          # if expression, dup the value we're assigning
          @method.dup if expression
          
          type.store(@method, @method.local(scoped_local_name(name, scope), type))
        end
        
        def declared_locals
          @declared_locals ||= {}
        end
        
        def annotate(builder, annotations)
          annotations.each do |annotation|
            type = annotation.type
            type = type.jvm_type if type.respond_to?(:jvm_type)
            builder.annotate(type, annotation.runtime?) do |visitor|
              annotation.values.each do |name, value|
                annotation_value(visitor, name, value)
              end
            end
          end
        end
        
        def annotation_value(builder, name, value)
          case value
          when Mirah::AST::Annotation
            type = value.type
            type = type.jvm_type if type.respond_to?(:jvm_type)
            builder.annotation(name, type) do |child|
              value.values.each do |name, value|
                annotation_value(child, name, value)
              end
            end
          when ::Array
            builder.array(name) do |array|
              value.each do |item|
                annotation_value(array, nil, item)
              end
            end
          else
            builder.value(name, value)
          end
        end
        
        def declared?(scope, name)
          declared_locals.include?(scoped_local_name(name, scope))
        end
        
        def declare_local(scope, name, type)
          # TODO confirm types are compatible
          name = scoped_local_name(name, scope)
          unless declared_locals[name]
            declared_locals[name] = type
            index = @method.local(name, type)
          end
        end
        
        def local_declare(scope, name, type)
        end
        
        def declare_locals(scope)
          scope.locals.each do |name|
            unless scope.captured?(name) || declared?(scope, name)
              type = scope.local_type(name)
              declare_local(scope, name, type)
              type.init_value(@method)
              type.store(@method, @method.local(scoped_local_name(name, scope), type))
            end
          end
        end
        
        def get_binding(type)
          @bindings[type]
        end
        
        def declared_captures(binding=nil)
          @captured_locals[binding || @binding]
        end
        
        def captured_local_declare(scope, name, type)
          unless declared_captures[name]
            declared_captures[name] = type
            # default should be fine, but I don't think bitescript supports it.
            @binding.protected_field(name, type)
          end
        end
        
        def captured_local(scope, name, type)
          captured_local_declare(scope, name, type)
          binding_reference
          @method.getfield(scope.binding_type, name, type)
        end
        
        def captured_local_assign(node, expression)
          scope, name, type = node.containing_scope, node.name, node.inferred_type
          captured_local_declare(scope, name, type)
          binding_reference
          node.value.compile(self, true)
          @method.dup_x2 if expression
          @method.putfield(scope.binding_type, name, type)
        end
        
        def field(name, type, annotations, static_field)
          name = name[1..-1] if name =~ /^@/
          
          real_type = declared_fields[name] || type
          
          declare_field(name, real_type, annotations, static_field)
          
          # load self object unless static
          method.aload 0 unless static || static_field
          
          if static || static_field
            @method.getstatic(@class, name, type)
          else
            @method.getfield(@class, name, type)
          end
        end
        
        def declared_fields
          @declared_fields ||= {}
          @declared_fields[@class] ||= {}
        end
        
        def declare_field(name, type, annotations, static_field)
          # TODO confirm types are compatible
          unless declared_fields[name]
            declared_fields[name] = type
            field = if static || static_field
              @class.private_static_field name, type
            else
              @class.private_field name, type
            end
            annotate(field, annotations)
          end
        end
        
        def field_declare(name, type, annotations, static_field)
          name = name[1..-1] if name =~ /^@/
          declare_field(name, type, annotations, static_field)
        end
        
        def field_assign(name, type, expression, value, annotations, static_field)
          name = name[1..-1] if name =~ /^@/
          
          real_type = declared_fields[name] || type
          
          declare_field(name, real_type, annotations, static_field)
          
          method.aload 0 unless static || static_field
          value.compile(self, true)
          if expression
            instruction = 'dup'
            instruction << '2' if type.wide?
            instruction << '_x1' unless static || static_field
            method.send instruction
          end
          
          if static || static_field
            @method.putstatic(@class, name, real_type)
          else
            @method.putfield(@class, name, real_type)
          end
        end
        
        def string(value)
          @method.ldc(value)
        end
        
        def build_string(nodes, expression)
          if expression
            # could probably be more efficient with non-default constructor
            builder_class = Mirah::AST.type(nil, 'java.lang.StringBuilder')
            @method.new builder_class
            @method.dup
            @method.invokespecial builder_class, "<init>", [@method.void]
            
            nodes.each do |node|
              node.compile(self, true)
              method = find_method(builder_class, "append", [node.inferred_type], false)
              if method
                @method.invokevirtual builder_class, "append", [method.return_type, *method.argument_types]
              else
                log "Could not find a match for #{java::lang::StringBuilder}.append(#{node.inferred_type})"
                fail "Could not compile"
              end
            end
            
            # convert to string
            @method.invokevirtual java::lang::StringBuilder.java_class, "toString", [@method.string]
          else
            nodes.each do |node|
              node.compile(self, false)
            end
          end
        end
        
        def to_string(body, expression)
          if expression
            body.compile(self, true)
            body.inferred_type.box(@method) if body.inferred_type.primitive?
            null = method.label
            done = method.label
            method.dup
            method.ifnull(null)
            @method.invokevirtual @method.object, "toString", [@method.string]
            @method.goto(done)
            null.set!
            method.pop
            method.ldc("null")
            done.set!
          else
            body.compile(self, false)
          end
        end
        
        def boolean(value)
          value ? @method.iconst_1 : @method.iconst_0
        end
        
        def regexp(value, flags = 0)
          # TODO: translate flags to Java-appropriate values
          @method.ldc(value)
          @method.invokestatic java::util::regex::Pattern, "compile", [java::util::regex::Pattern, @method.string]
        end
        
        def array(node, expression)
          if expression
            # create basic arraylist
            @method.new java::util::ArrayList
            @method.dup
            @method.ldc_int node.children ? node.children.size : 0
            @method.invokespecial java::util::ArrayList, "<init>", [@method.void, @method.int]
            
            # elements, as expressions
            # TODO: ensure they're all reference types!
            node.children.each do |n|
              @method.dup
              n.compile(self, true)
              # TODO this feels like it should be in the node.compile itself
              if n.inferred_type!.primitive?
                n.inferred_type.box(@method)
              end
              @method.invokeinterface java::util::List, "add", [@method.boolean, @method.object]
              @method.pop
            end
            
            # make it unmodifiable
            @method.invokestatic java::util::Collections, "unmodifiableList", [java::util::List, java::util::List]
          else
            # elements, as non-expressions
            # TODO: ensure they're all reference types!
            node.children.each do |n|
              n.compile(self, true)
              # TODO this feels like it should be in the node.compile itself
              if n.inferred_type.primitive?
                n.inferred_type.box(@method)
              end
            end
          end
        end
        
        def null
          @method.aconst_null
        end
        
        def binding_reference
          @method.aload(@method.local('$binding'))
        end
        
        def real_self
          method.aload(0)
        end
        
        def line(num)
          @method.line(num - 1) if @method
        end
        
        def print(print_node)
          @method.getstatic System, "out", PrintStream
          print_node.parameters.each {|param| param.compile(self, true)}
          params = print_node.parameters.map {|param| param.inferred_type.jvm_type}
          method_name = print_node.println ? "println" : "print"
          method = find_method(PrintStream.java_class, method_name, params, false)
          if (method)
            @method.invokevirtual(
            PrintStream,
            method_name,
            [method.return_type, *method.parameter_types])
          else
            log "Could not find a match for #{PrintStream}.#{method_name}(#{params})"
            fail "Could not compile"
          end
        end
        
        def return(return_node)
          return_node.value.compile(self, true) if return_node.value
          handle_ensures(find_ensures(Mirah::AST::MethodDefinition))
          return_node.inferred_type.return(@method)
        end
        
        def _raise(exception)
          exception.compile(self, true)
          @method.athrow
        end
        
        def rescue(rescue_node, expression)
          start = @method.label.set!
          body_end = @method.label
          done = @method.label
          rescue_node.body.compile(self, expression && rescue_node.else_node.nil?) if rescue_node.body
          body_end.set!
          rescue_node.else_node.compile(self, expression) if rescue_node.else_node
          return if start.label.offset == body_end.label.offset
          @method.goto(done)
          rescue_node.clauses.each do |clause|
            target = @method.label.set!
            if clause.name
              @method.astore(declare_local(clause.static_scope, clause.name, clause.type))
            else
              @method.pop
            end
            declare_locals(clause.static_scope)
            clause.body.compile(self, expression)
            @method.goto(done)
            clause.types.each do |type|
              @method.trycatch(start, body_end, target, type)
            end
          end
          done.set!
        end
        
        def handle_ensures(nodes)
          nodes.each do |ensure_node|
            ensure_node.clause.compile(self, false)
          end
        end
        
        def ensure(node, expression)
          node.state = @method.label  # Save the ensure target for JumpNodes
          start = @method.label.set!
          body_end = @method.label
          done = @method.label
          push_jump_scope(node) do
            node.body.compile(self, expression)  # First compile the body
          end
          body_end.set!
          handle_ensures([node])  # run the ensure clause
          @method.goto(done)  # and continue on after the exception handler
          target = @method.label.set!  # Finally, create the exception handler
          @method.trycatch(start, body_end, target, nil)
          handle_ensures([node])
          @method.athrow
          done.set!
        end
        
        def empty_array(type, size)
          size.compile(self, true)
          type.newarray(@method)
        end
        
        class ClosureCompiler < JVMBytecode
          def initialize(file, type, parent)
            @file = file
            @type = type
            @jump_scope = []
            @parent = parent
          end
          
          def prepare_binding(scope)
            if scope.has_binding?
              type = scope.binding_type
              @binding = @parent.get_binding(type)
              @method.aload 0
              @method.getfield(@class, 'binding', @binding)
              type.store(@method, @method.local('$binding', type))
            end
            begin
              yield
            ensure
              if scope.has_binding?
                @binding = nil
              end
            end
          end
          
          def declared_captures
            @parent.declared_captures(@binding)
          end
        end
      end
    end
  end
end