class Cucumber::Formatter::AstLookup

def gherkin_document(uri)
  @gherkin_documents[uri]
end

def initialize(config)
  @gherkin_documents = {}
  @test_case_lookups = {}
  @test_step_lookups = {}
  @step_keyword_lookups = {}
  config.on_event :gherkin_source_parsed, &method(:on_gherkin_source_parsed)
end

def on_gherkin_source_parsed(event)
  @gherkin_documents[event.gherkin_document.uri] = event.gherkin_document
end

def scenario_source(test_case)
  uri = test_case.location.file
  @test_case_lookups[uri] ||= TestCaseLookupBuilder.new(gherkin_document(uri)).lookup_hash
  @test_case_lookups[uri][test_case.location.lines.max]
end

def snippet_step_keyword(test_step)
  uri = test_step.location.file
  document = gherkin_document(uri)
  dialect = ::Gherkin::Dialect.for(document.feature.language)
  given_when_then_keywords = [dialect.given_keywords, dialect.when_keywords, dialect.then_keywords].flatten.uniq.reject { |kw| kw == '* ' }
  keyword_lookup = step_keyword_lookup(uri)
  keyword = nil
  node = keyword_lookup[test_step.location.lines.min]
  while keyword.nil?
    if given_when_then_keywords.include?(node.keyword)
      keyword = node.keyword
      break
    end
    break if node.previous_node.nil?
    node = node.previous_node
  end
  keyword = dialect.given_keywords.reject { |kw| kw == '* ' }[0] if keyword.nil?
  Cucumber::Gherkin::I18n.code_keyword_for(keyword)
end

def step_keyword_lookup(uri)
  @step_keyword_lookups[uri] ||= KeywordLookupBuilder.new(gherkin_document(uri)).lookup_hash
end

def step_source(test_step)
  uri = test_step.location.file
  @test_step_lookups[uri] ||= TestStepLookupBuilder.new(gherkin_document(uri)).lookup_hash
  @test_step_lookups[uri][test_step.location.lines.min]
end