module LNR

using Compat
using Lazy
import Base: seek, position, eof, peek

export LineNumberingReader, line, column, Cursor, cursor, seekline, seekcol

# Cursor Type

immutable Cursor
  line::Int
  column::Int
end

Base.isless(x::Cursor, y::Cursor) =
  x.line < y.line || (x.line == y.line && x.column < y.column)

# LNR

immutable LineNumberingReader{T<:IO} <: IO
  io::T
  lines::Vector{Int} # The byte index of the first char in each line
end

LineNumberingReader(io::IO) = LineNumberingReader(io, [1])
LineNumberingReader(s::AbstractString) = LineNumberingReader(IOBuffer(s))

eof(r::LineNumberingReader) = eof(r.io)
position(r::LineNumberingReader) = position(r.io)
peek(r::LineNumberingReader) = Base.peek(r.io)

scannedindex(r::LineNumberingReader, i) = i < r.lines[end]

function Base.read(r::LineNumberingReader, ::Type{UInt8})
  c = read(r.io, UInt8)
  c == '\n' && !scannedindex(r, position(r)) && !eof(r) &&
     push!(r.lines, position(r)+1)
  return c
end

function Base.skip(r::LineNumberingReader, n::Integer)
  if n > 0 && !scannedindex(r, position(r) + n)
    @dotimes n read(r, UInt8)
  else
    skip(r.io, n)
  end
end

# Seeking

function seek(r::LineNumberingReader, pos)
  scannedindex(r, pos) ? seek(r.io, pos) : skip(r, pos-position(r))
  return r
end

function seekline(r::LineNumberingReader, line::Integer)
  line ≤ length(r.lines) && return seek(r, r.lines[line]-1)
  while length(r.lines) < line && !eof(r)
    readline(r)
  end
  return r
end

seekline(r::LineNumberingReader) = seekline(r, line(r))

function seekcol(r::LineNumberingReader, col::Integer)
  seekline(r)
  c = 1
  while c < col && !eof(r)
    read(r, Char) == '\n' && return skip(r,-1)
    c += 1
  end
  return r
end

seek(r::LineNumberingReader, c::Cursor) =
  @> r seekline(c.line) seekcol(c.column)

# Cursor finding

function line(s::LineNumberingReader)
  p = position(s)+1 # after the cursor
  for i = 1:length(s.lines)
    s.lines[i] > p && return i - 1
  end
  return length(s.lines)
end

function column(r::LineNumberingReader)
  pos = position(r)
  col = 1
  seekline(r)
  while position(r) < pos
    read(r, Char)
    col += 1
  end
  return col
end

cursor(s::LineNumberingReader) = Cursor(line(s), column(s))
cursor(l, c) = Cursor(l, c)

function position(r::LineNumberingReader, p::Cursor)
  withstream(r) do
    seek(r, p)
    position(r)
  end
end

# Position is 0-based
index(r::LineNumberingReader, p::Cursor) = position(r, p) + 1

function peekchar(r::LineNumberingReader)
  withstream(r) do
    read(r, Char)
  end
end

# Util

function withstream(f, io)
  pos = position(io)
  try
    f()
  finally
    seek(io, pos)
  end
end

end
