require_relative "./ops.rb"
require_relative "./lex.rb"

AST = Struct.new(:token, :args, :type, :impl, :label, :is_fold_op)

def split_ids(tokens, registers)
  tokens.map{|t|
    if t.type==:id && !lookup_op(t.str) && !registers[t.str]
      warn("splitting %p into chars because no op/var by that name, but it is surrounded by whitespace" % t.str, t) if t.surrounded_by_whitespace
      t.split
    else
      t
    end
  }.flatten
end

Keywords = ["set","let",">","=","del"]
def get_var_ids_set(tokens)
  names = {}
  tokens.size.times{|i|
    if tokens[i].str == "set" || tokens[i].str == "let"
      raise IogiiError.new "an id must follow `#{tokens[i].str}`", tokens[i] if !tokens[i+1] || Keywords.include?(tokens[i+1].str)
      names[tokens[i+1].str] = true
    end
  }

  nregisters = tokens.filter{|t|t.str == "="}.size
  used = [false]*26
  tokens.filter{|t|t.type == :id }.map(&:str).join.scan(/[A-Z]/){|s|
    used[s.ord-?A.ord] = true }
  starts = (1...26).filter{|i| used[i] && !used[i-1] }.unshift 0
  diffs = starts.map.with_index{|s,i| (starts[i+1]||26)-starts[i] }
  largest_diff_at = starts[diffs.index(diffs.max)]
  consec_used = used[largest_diff_at..-1].take_while{|i| i }.size
  auto_names = (0...nregisters).map{|x|
    (x + ?A.ord + largest_diff_at + [consec_used - nregisters,0].max).chr}

  auto_names.each{|s| names[s] = true }
  tokens.filter{|t|t.str == "="}.each.with_index{|t,i|
    name = auto_names[i]
    t.sets_var = name
    raise IogiiError.new"Out of letters for auto vars", t if name > ?Z
    raise IogiiError.new"Sets register '#{name}' but it is never used", t if !used[name.ord-?A.ord]
  }

  names
end

def is_datum(token)
  [:int,:char,:str].include?(token.type)
end

def parse_data(tokens, token_ind)
  i = token_ind - 1
  data_tokens = []
  prev = false
  while i < tokens.size
    if is_datum(tokens[i])
      break if prev
      if tokens.size > i+1 && tokens[i].type == :int && tokens[i+1].str[0] == "~" && tokens[i+1].str[-1] == ","
        data_tokens << tokens[i].dup
        data_tokens << tokens[i+1].dup
        data_tokens[-2].str = "-"+data_tokens[-2].str
        data_tokens[-1].str = data_tokens[-1].str[1..-1]
        data_tokens[-1].type == :commas
        i += 1
        prev = false
      else
        prev = true
        data_tokens << tokens[i]
      end
    else
      break if tokens[i].type != :commas
      prev = false
      data_tokens << tokens[i]
    end
    i += 1
  end
  tokens = data_tokens

  depth = tokens.reject{|t|is_datum(t)}.map{|t|t.str.size}.max||0
  token_type = nil
  tokens.select{|t|is_datum(t)}.each{|t|
    token_type ||= t.type
    token_type = :str if token_type != t.type
  }

  value = parse_list(tokens,depth,token_type)
  case token_type
  when :str
    type = CharType
    depth += 1
  when :char
    type = CharType
  when :int
    type = IntType
  else; raise
  end
  [new_op("data", type_and_rank_to_str(type,depth)){ value }, i]
end

def split_tokens(tokens, str)
  tokens.slice_when{|t| t.str == str }.
    map{|ts| ts.reject{|t| t.str == str } }
end

def parse_list(tokens,depth,to_type)
  if depth == 0
    raise IogiiError.new 'found ", ," invalid data format' if tokens.empty?
    parse_datum(tokens[0],to_type)
  else
    to_lazy_list(split_tokens(tokens, ","*depth).map{|ts|
      parse_list(ts, depth-1, to_type)
    })
  end
end

def parse_datum(token,to_type)
  case token.type
  when :int
    val = token.str.to_i
    to_type == :str ? str(val.const) : val
  when :str
    str,remainder = *parse_str(token.str)
    if remainder == '"b' && str.bytes.all?{|b| b < 128 }
      warn("You have used the raw byte mode on a string that does not contain any binary characters, if you wish to reverse the string instead then add a space before the `b`", token)
    end
    str_to_lazy_list(str)
  when :char
    raise IogiiError.new "empty char", token if token.str.size < 2
    x = parse_char(token.str[1..-1]).ord
    to_type == :str ? [x.const, Null] : x
  end
end

def tokenize_data(tokens, registers)
  new_tokens = []
  token_ind = 0
  loop {
    t = tokens[token_ind]
    token_ind += 1
    break if token_ind > tokens.size
    if registers[t.str] || !is_datum(t)
      new_tokens << t
    else
      orig_token_ind = token_ind
      orig_token = tokens[token_ind-1]
      data_impl, token_ind = *parse_data(tokens, token_ind)
      str = tokens[orig_token_ind-1...token_ind].map(&:str).join
      new_token = Token.new(str, :data, orig_token.char_no, orig_token.line_no)
      new_token.data_impl = data_impl
      new_tokens << new_token
    end
  }
  new_tokens
end

# todo, this is a more expensive of an operation than it needs to be
def is_ancestor(parent, child, nodes)
  return true if parent == child
  nodes[parent].args.any?{|arg| is_ancestor(arg, child, nodes) }
end

def route_to_ancestor(parent, child, nodes)
  return [parent] if parent == child
  nodes[parent].args.each{|arg|
    route = route_to_ancestor(arg, child, nodes)
    return route << parent if route
  }
  nil
end

def build_tree(tokens, registers, input_present, last_unmatched = nil, first_right_curly_brace = nil)
  # it is important for each token to only be responsible for adding 1 value to the stack, that way rotating the main program can choose a break point at any point. That's why adding multiple is accomplished by creating virtual tokens that are responsible for the other values.

  stack = []
  lbs = []
  loop_lbs = []
  all_lbs = []
  nodes = []
  deleted = []
  token_ind = 0
  nimplicits = 0
  t = blocks_can_close_now = nil
  original_last_unmatched = last_unmatched
  original_first_right_curly_brace = first_right_curly_brace
  first_right_curly_brace = nil

  push = -> a { stack << nodes.size; nodes << a; a }
  pop = -> {
    if stack.empty?
      nimplicits += 1
      push[AST.new(t, [], :ImplicitValue)]
    end
    if nodes[stack.last].type == :ImplicitValue
      nodes[stack.last].impl = if last_unmatched
        # only set this on final pass, since one of those implicits could become input
        last_unmatched.token.implicit_used = true if original_last_unmatched
        last_unmatched.args[0]
      elsif all_lbs.empty? && !input_present.value && original_first_right_curly_brace
        original_first_right_curly_brace
      else
        all_lbs[-1]
      end
    end
    stack.pop
  }
  peek = -> { ans = pop[]; stack << ans; ans }
  # add a token for this so that it can be wrapped
  insert_duped = -> dup_ind,pos {
    t_dup = nodes[dup_ind].token
    if !t_dup.duped_value # (first pass doesn't have duped tokens yet)
      t_dup.duped_value = pos.map{|i|
        tokens[token_ind,0] = Token.new(i == 0 ? "duped_value" : "peek_placeholder", :"VirtualDuped#{i}")
      }
    end
    t_dup.duped_value.each{|v|v.duped_value = nodes[dup_ind] }
    blocks_can_close_now = false
  }

  loop {
    t = tokens[token_ind]
    token_ind += 1
    blocks_can_close_now = true
    if token_ind > tokens.size
      break
    elsif registers[t.str]
      push[AST.new(t, [], :VarNode)]
    elsif t.str && t.str[0] == ">"
      if t.matching_bracket
        lb = nodes.size.times.find{|node| nodes[node].token.equal?(t.matching_bracket) }
        all_lbs.delete(lb)
        push[AST.new(t, [pop[]], :MatchedEndBracket, lb)]
        nodes[lb].impl = nodes.size-1
      else
        push[last_unmatched = AST.new(t, [pop[]], :UnmatchedEndBracket)]
      end
    elsif t.str == "set" || t.str == "let"
      name = tokens[token_ind].str
      token_ind += 1
      registers[name] = peek[]
      push[AST.new(t, [pop[]], :IdentityNode, nil, "set #{name}")]
      deleted << pop[] if t.str == "let"
    elsif t.str == "del"
      deleted << pop[]
    elsif t.str == "="
      registers[t.sets_var] = peek[]
      push[AST.new(t, [pop[]], :IdentityNode, nil, "set #{t.sets_var}")]
    elsif t.type == :data
      push[AST.new(t, [], :DataNode)]
    elsif t.type == :commas
      raise IogiiError.new "commas must follow data or op that can vectorize", t
    elsif t.type == :VirtualDuped0 || t.type == :VirtualDuped1
      push[AST.new(t, [], t.type)]
    elsif t.type == :id || t.type == :sym
      op = lookup_op(t.str)
      raise IogiiError.new "unknown op", t if !op
      args = op.parse_nargs.times.map{ pop[] }.reverse
      if op.name == "dup"
        push[AST.new(t, args, :IdentityNode)]
        insert_duped[nodes.size-1,[0]]
      elsif op.name == "peek"
        push[AST.new(t, args, :IdentityNode)]
        insert_duped[nodes.size-1,[0,1]]
      elsif op.name == "$"
        if !all_lbs.empty?
          push[AST.new(t, [], :ArgNode, all_lbs[-1])]
        elsif input_present.value
          nodes << mk_input_node
          push[AST.new(t, [], :ArgNode, nodes.size - 1)]
        else
          push[AST.new(t, [], :ImplicitValue)]
        end
      elsif op.name == "input"
        push[mk_input_node(t)]
      elsif op.name == "implicit"
        push[AST.new(t, [], :ImplicitValue)]
      elsif op.block
        if ['mdup', 'mpeek'].include?(op.name)
          push[AST.new(t, args, :MDupOp)]
          lbs << nodes.size - 1
          blocks_can_close_now = false
        else
          n = AST.new(t, args, :LoopOp)
          n.is_fold_op = ["foldr","foldr0"].include? op.name
          push[n]
          loop_lbs << nodes.size - 1 if !t.matching_bracket
        end
        all_lbs << nodes.size - 1
      else
        push[AST.new(t, args, :OpNode)]
        if t.str == "}" && !first_right_curly_brace
          first_right_curly_brace = nodes.size - 1
        end
      end
    else
      raise "unknown token type %p" % t.type.to_s
    end

    # implicitly close meta ops
    while blocks_can_close_now &&
        !lbs.empty? &&
        is_ancestor(peek[], lbs.last, nodes)
      lb = nodes[lb_ind = lbs.pop]
      all_lbs.delete(lb_ind)
      push[AST.new(t, [pop[]], :MatchedEndBracket, lb_ind)]
      lb.impl = nodes.size-1
      insert_duped[lb_ind,[0]] if ["mdup",";"].include?(lb.token.str)
      insert_duped[lb_ind,[0,1]] if ["mpeek","!"].include?(lb.token.str)
    end

    # implicitly close loops
    while blocks_can_close_now &&
        !loop_lbs.empty? &&
        peek[] != loop_lbs[-1] && # since loop ops don't stop blocks from closing, check to make sure there has been at least one op
        is_ancestor(peek[], loop_lbs.last, nodes)
      break if nodes[loop_lbs.last].is_fold_op && !route_to_ancestor(peek[], loop_lbs.last, nodes).any?{|n| nodes[n].args.size > 1 }
      lb = nodes[lb_ind = loop_lbs.pop]
      all_lbs.delete(lb_ind)
      push[AST.new(t, [pop[]], :MatchedEndBracket, lb_ind)]
      lb.impl = nodes.size-1
    end
  }

  # touch each element in stack to fill in implicit data
  stack = stack.size.times.map{ pop[] }.reverse

  [nodes, stack, deleted, nimplicits, all_lbs, first_right_curly_brace]
end

def mk_input_node(token = nil)
  $input_checked = true
  AST.new(token, [], :InputNode)
end

def match_brackets(nodes, root)
  known_non_matchers = {}

  # repeat this handling one `>,` at a time
  loop {
    loop_ops = [] # list of loop op inds (left to right)
    end_ops = [] # list of end brackets (left to right)
    loop_options = [] # list of end bracket inds loop op of same ind could match (right to left)
    end_options = [] # left to right

    get_loop_ops = -> node, end_options {
      is_unmatched_end_op = nodes[node].type == :UnmatchedEndBracket && !known_non_matchers[node]
      sub_end_options = is_unmatched_end_op ? end_options + [node] : end_options
      nodes[node].args.each{|i| get_loop_ops[i, sub_end_options] }
      if nodes[node].type == :LoopOp && !nodes[node].token.matching_bracket
        loop_ops << node
        loop_options << end_options
      elsif is_unmatched_end_op
        end_ops << node
      end
    }
    get_loop_ops[root, []]
    end_ops.each{|end_op|
      end_options << loop_ops.values_at(*loop_options.size.times.select{|i| loop_options[i].include?(end_op) })
    }

    end_matches = {} # end bracket ind -> loop ind
    loop_matches = {} # loop ind -> end bracket ind

    # greedy first pass to figure out which loop ops are matched (minimal match length though)
    # this guarantees maximum number of matches found
    (loop_ops.size-1).downto(0){|i|
      loop_op = loop_ops[i]
      option = loop_options[i].reverse.find{|e| !end_matches[e] }
      if option
        end_matches[option] = loop_op
        loop_matches[loop_op] = option
      end
    }

    progress = true
    while progress
      progress = false
      # extend each loop op match to a further end bracket when possible
      # does not cause crossings because it starts with leftmost loop op,
      # if a later loop op could extend crossing and earlier one, the earlier one would already have since it is known to have closed between two two points
      loop_ops.size.times{|i|
        loop_op = loop_ops[i]
        option = loop_options[i].find{|e| !end_matches[e] }
        if option && option > loop_matches[loop_op]
          progress = true
          end_matches.delete(loop_matches[loop_op])
          end_matches[option] = loop_op
          loop_matches[loop_op] = option
        end
      }

      # extend each end bracket to match an earlier loop op where possible
      # does not cause crossings because it starts with rightmost end bracket
      end_ops.size.times.reverse_each{|i|
        end_op = end_ops[i]
        option = end_options[i].find{|b| !loop_matches[b] }
        if option && option < end_matches[end_op]
          progress = true
          loop_matches.delete(end_matches[end_op])
          loop_matches[option] = end_op
          end_matches[end_op] = option
        end
      }
    end

    last_unmatched_comma_end_bracket = end_ops.reverse.find{|end_ind|
      end_op = nodes[end_ind]
      end_op.type == :UnmatchedEndBracket && end_op.token.str[","] && !known_non_matchers[end_ind]
    }
    if last_unmatched_comma_end_bracket
      end_node = nodes[last_unmatched_comma_end_bracket]
      comma_count = end_node.token.str.count(",")
      current_match = end_matches[last_unmatched_comma_end_bracket]
      if !current_match
        raise IogiiError.new "found >, but this end bracket already didn't match a loop op so that `,` is unnecessary", end_node
      end
      end_op = end_ops.index(last_unmatched_comma_end_bracket)
      new_options = end_options[end_op][end_options[end_op].index(current_match) + 1..-1]
      if comma_count > new_options.size # set implicit
        known_non_matchers[last_unmatched_comma_end_bracket] = true
      else # move match
        new_match = new_options[comma_count-1]
        end_node.type = :MatchedEndBracket
        end_node.token.matching_bracket = nodes[new_match].token
        nodes[new_match].token.matching_bracket = end_node.token
      end
    else # no more comma end brackets to handle, this is final matching configuration
      loop_matches.each{|loop_ind,end_ind|
        nodes[end_ind].type = :MatchedEndBracket
        nodes[end_ind].token.matching_bracket = nodes[loop_ind].token
        nodes[loop_ind].token.matching_bracket = nodes[end_ind].token
      }
      break
    end
  }
end

def parse_main(tokens, input_present)
  registers = get_var_ids_set(tokens)
  tokens = split_ids(tokens, registers)
  tokens = tokenize_data(tokens, registers)
  nodes, stack, deleted, nimplicits, lbs, first_right_curly_brace = build_tree(tokens, registers, input_present)
  if nimplicits > 0
    unfinished_extra = lbs.map{|lb|
      case nodes[lb].token.str
        when ";", "mdup"; 1
        when "!", "mpeek"; 2
        else; 0
      end
    }.sum
    # we don't need as many implicits since rotating will add extra duped values not yet present
    nimplicits = [nimplicits - unfinished_extra, 1].max
    nrot = [nimplicits - 1, stack.size - 1].min

    if stack.empty?
      last_token_ind = tokens.size
    else
      last_token = nodes[stack[~nrot]].token
      last_token_ind = tokens.rindex{|t|t.equal? last_token}
      last_token_ind += last_token.str == "let" || last_token.str == "set" ? 2 : 1
    end

    input = Token.new(input_present.value ? "input" : "implicit", :id)

    tokens = tokens[last_token_ind..-1] + [input] + tokens[0,last_token_ind]

    # build the tree after literally rotating the tokens, rather than manipulate the tree because the program rotation could have interupted functions which would not have been able to implicitly close
    nodes, stack, deleted, nimplicits, lbs, first_right_curly_brace = build_tree(tokens, registers, input_present)
  end

#     puts tokens.map(&:str)*" "
# puts
  (stack + deleted).each{|si| match_brackets(nodes, si) }

  last_unmatched = nodes.reverse_each.find{|n| n.type == :UnmatchedEndBracket }

  nodes, stack, deleted, nimplicits, lbs, first_right_curly_brace = build_tree(tokens, registers, input_present, last_unmatched, first_right_curly_brace)

  if last_unmatched || first_right_curly_brace
    # this is a bit hacky, but do it again since last_unmatched is an index to the ast which can change from when it was calculated. example program that requires this:
    # i>e>5>
    last_unmatched = nodes.reverse_each.find{|n| n.type == :UnmatchedEndBracket }
    nodes, stack, deleted, nimplicits, lbs, first_right_curly_brace = build_tree(tokens, registers, input_present, last_unmatched, first_right_curly_brace)
  end

  raise IogiiError.new("unterminated function (stack is never rid of extra values again, functions can only return one value)\n" + draw_fns(tokens, nodes), nodes[lbs[-1]]) if !lbs.empty?

  if last_unmatched
    nodes.each{|n|
      raise IogiiError.new "unmatched bracket // implicit value is set but not used\n" + draw_fns(tokens, nodes), n if n.type == :UnmatchedEndBracket && !n.token.implicit_used
    }
  else
    nodes.filter{|n| n.type == :ImplicitValue }.each{|n|
      if n.impl == nil
        nodes << mk_input_node
        n.impl = nodes.size - 1
      end
    }
  end

  [nodes, main_inds=stack, registers]
end
