-----------------------------------------------------------------------------
-- Little program that checks links in HTML files, using coroutines and
-- non-blocking I/O via the dispatcher module.
-- LuaSocket sample files
-- Author: Diego Nehab
-----------------------------------------------------------------------------
local url = require("socket.url")
local dispatch = require("dispatch")
local http = require("socket.http")
dispatch.TIMEOUT = 10

-- make sure the user knows how to invoke us
arg = arg or {}
if #arg < 1 then
    print("Usage:\n  luasocket check-links.lua [-n] {<url>}")
    exit()
end

-- '-n' means we are running in non-blocking mode
if arg[1] == "-n" then
    -- if non-blocking I/O was requested, use real dispatcher interface
    table.remove(arg, 1)
    handler = dispatch.newhandler("coroutine")
else
    -- if using blocking I/O, use fake dispatcher interface
    handler = dispatch.newhandler("sequential")
end

local nthreads = 0

-- get the status of a URL using the dispatcher
function getstatus(link)
    local parsed = url.parse(link, {scheme = "file"})
    if parsed.scheme == "http" then
        nthreads = nthreads + 1
        handler:start(function()
            local r, c, h, s = http.request{
                method = "HEAD",
                url = link,
                create = handler.tcp
            }
            if r and c == 200 then io.write('\t', link, '\n')
            else io.write('\t', link, ': ', tostring(c), '\n') end
            nthreads = nthreads - 1
        end)
    end
end

function readfile(path)
    path = url.unescape(path)
    local file, error = io.open(path, "r")
    if file then
        local body = file:read("*a")
        file:close()
        return body
    else return nil, error end
end

function load(u)
    local parsed = url.parse(u, { scheme = "file" })
    local body, headers, code, error
    local base = u
    if parsed.scheme == "http" then
        body, code, headers = http.request(u)
        if code == 200 then
            -- if there was a redirect, update base to reflect it
            base = headers.location or base
        end
        if not body then
            error = code
        end
    elseif parsed.scheme == "file" then
        body, error = readfile(parsed.path)
    else error = string.format("unhandled scheme '%s'", parsed.scheme) end
    return base, body, error
end

function getlinks(body, base)
    -- get rid of comments
    body = string.gsub(body, "%<%!%-%-.-%-%-%>", "")
    local links = {}
    -- extract links
    body = string.gsub(body, '[Hh][Rr][Ee][Ff]%s*=%s*"([^"]*)"', function(href)
        table.insert(links, url.absolute(base, href))
    end)
    body = string.gsub(body, "[Hh][Rr][Ee][Ff]%s*=%s*'([^']*)'", function(href)
        table.insert(links, url.absolute(base, href))
    end)
    string.gsub(body, "[Hh][Rr][Ee][Ff]%s*=%s*(.-)>", function(href)
        table.insert(links, url.absolute(base, href))
    end)
    return links
end

function checklinks(address)
    local base, body, error = load(address)
    if not body then print(error) return end
    print("Checking ", base)
    local links = getlinks(body, base)
    for _, link in ipairs(links) do
        getstatus(link)
    end
end

for _, address in ipairs(arg) do
    checklinks(url.absolute("file:", address))
end

while nthreads > 0 do
    handler:step()
end