#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)
# Copyright (C) Guanlan Dai

# TODO: port this script into the nginx core for greater flexibility
# and better performance.

use strict;
use warnings;

our $VERSION = '0.13';

use File::Spec ();
use FindBin ();
use List::Util qw( max );
use Getopt::Long qw( GetOptions :config no_ignore_case require_order);
use File::Temp qw( tempdir );
use POSIX qw( WNOHANG );

sub resolve_includes {
    my $type = shift;
    my @paths = map {
        my $abs_path = File::Spec->rel2abs($_);
        unless (-f $abs_path) {
          die "Could not find $type include '$abs_path'";
        }
        "include $abs_path;";
    } @_;
    return join("\n", @paths);
}

my @all_args = @ARGV;

my (@http_includes, @main_includes, @src_a);

GetOptions("c=i",             \(my $conns_num),
           "e=s",             \@src_a,
           "h|help",          \(my $help),
           "http-include=s",  \@http_includes,
           "I=s@",            \(my $Inc),
           "main-include=s",  \@main_includes,
           "nginx=s",         \(my $nginx_path),
           "valgrind",        \(my $use_valgrind),
           "valgrind-opts=s", \(my $valgrind_opts),
           "V|v",             \(my $version))
   or die usage(1);

my $src;
if (@src_a) {
    $src = join('; ', @src_a);
}

if ($help) {
    usage(0);
}

if (!$nginx_path) {
    use Config;
    my $ext = $Config{_exe};
    if (!$ext) {
        if ($^O eq 'msys') {
            $ext = '.exe';
        } else {
            $ext = '';
        }
    }
    $nginx_path = File::Spec->catfile($FindBin::RealBin, "..", "nginx", "sbin", "nginx$ext");
    if (!-f $nginx_path) {
        $nginx_path = File::Spec->catfile($FindBin::RealBin, "nginx$ext");
        if (!-f $nginx_path) {
            $nginx_path = "nginx";  # find in PATH
        }
    }
}

#warn $nginx_path;

if ($version) {
    warn "resty $VERSION\n";
    my $cmd = "$nginx_path -V";
    exec $cmd or die "Failed to run command \"$cmd\": $!\n";
}

my $lua_package_path_config = '';
if (defined $Inc) {
    my $package_path = "";
    my $package_cpath = "";
    for my $dir (@$Inc) {
        if (!-d $dir) {
            die "Search directory $dir is not found.\n";
        }
        $package_path .= File::Spec->catfile($dir, "?.lua;");
        $package_cpath .= File::Spec->catfile($dir, "?.so;");
    }
    $lua_package_path_config = <<_EOC_;
    lua_package_path "$package_path;";
    lua_package_cpath "$package_cpath;";
_EOC_
}

my $luafile = shift;
if (!defined $src and !defined $luafile) {
    die qq{Neither Lua input file nor -e "" option specified.\n};
}

my $conns = $conns_num || 64;

my @nameservers;

# try to read the nameservers used by the system resolver:
if (open my $in, "/etc/resolv.conf") {
    while (<$in>) {
        if (/^\s*nameserver\s+(\d+(?:\.\d+){3})(?:\s+|$)/) {
            push @nameservers, $1;
            if (@nameservers > 10) {
                last;
            }
        }
    }
    close $in;
}

if (!@nameservers) {
    # default to Google's open DNS servers
    push @nameservers, "8.8.8.8", "8.8.4.4";
}

#warn "@nameservers\n";

my $prefix_dir;
if ($^O eq 'msys') {
    # to work around a bug in msys perl (at least 5.8.8 msys 64int)
    $prefix_dir = "resty_cli_temp";
    if (-d $prefix_dir) {
        system("rm -rf $prefix_dir") == 0 or die $!;
    }
    mkdir $prefix_dir or die "failed to mkdir $prefix_dir: $!";

} else {
    $prefix_dir = tempdir(CLEANUP => 1);
    if ($^O eq 'MSWin32') {
        require Win32;
        $prefix_dir = Win32::GetLongPathName($prefix_dir);
    }
}
#warn "prefix dir: $prefix_dir\n";

my $logs_dir = File::Spec->catfile($prefix_dir, "logs");
mkdir $logs_dir or die "failed to mkdir $logs_dir: $!";

my $conf_dir = File::Spec->catfile($prefix_dir, "conf");
mkdir $conf_dir or die "failed to mkdir $conf_dir: $!";

my $inline_lua = '';
my $quoted_luafile;
if (defined $src) {
    my $file = File::Spec->catfile($conf_dir, "a.lua");
    open my $out, ">$file"
        or die "Cannot open $file for writing: $!\n";
    print $out $src;
    close $out;
    my $chunk_name = "=(command line -e)";
    $quoted_luafile = quote_as_lua_str($file);

    $inline_lua = <<"_EOC_";
                local fname = $quoted_luafile
                local f = assert(io.open(fname, "r"))
                local chunk = f:read("*a")
                local inline_gen = assert(loadstring(chunk, "$chunk_name"))
_EOC_
}

my $file_lua = '';
if (defined $luafile) {
    my $chunk_name = "\@$luafile";
    $quoted_luafile = quote_as_lua_str($luafile);
    $file_lua = <<"_EOC_";
                local fname = $quoted_luafile
                local f = assert(io.open(fname, "r"))
                local chunk = f:read("*a")
                local file_gen = assert(loadstring(chunk, "$chunk_name"))
_EOC_
}

my @user_args =  @ARGV;
my $args = gen_lua_code_for_args(\@user_args, \@all_args);

my $loader = <<_EOC_;
            local gen
            do
                $args
$inline_lua
$file_lua

                gen = function()
                  if inline_gen then inline_gen() end
                  if file_gen then file_gen() end
                end
            end
_EOC_

my $env_list = '';
for my $var (sort keys %ENV) {
    #warn $var;
    $env_list .= "env $var;\n";
}

my $main_include_directives = resolve_includes('main', @main_includes);
my $http_include_directives = resolve_includes('http', @http_includes);

my $conf_file = File::Spec->catfile($conf_dir, "nginx.conf");
open my $out, ">$conf_file"
    or die "Cannot open $conf_file for writing: $!\n";

print $out <<_EOC_;
daemon off;
master_process off;
worker_processes 1;
pid logs/nginx.pid;

$env_list

error_log stderr warn;
#error_log stderr debug;

events {
    worker_connections $conns;
}

$main_include_directives

http {
    access_log off;
    lua_socket_log_errors off;
    resolver @nameservers;
$lua_package_path_config
    $http_include_directives
    init_by_lua '
        local stdout = io.stdout
        local ngx_null = ngx.null
        local maxn = table.maxn
        local unpack = unpack
        local concat = table.concat

        local expand_table
        function expand_table(src, inplace)
            local n = maxn(src)
            local dst = inplace and src or {}
            for i = 1, n do
                local arg = src[i]
                local typ = type(arg)
                if arg == nil then
                    dst[i] = "nil"

                elseif typ == "boolean" then
                    if arg then
                        dst[i] = "true"
                    else
                        dst[i] = "false"
                    end

                elseif arg == ngx_null then
                    dst[i] = "null"

                elseif typ == "table" then
                    dst[i] = expand_table(arg, false)

                elseif typ ~= "string" then
                    dst[i] = tostring(arg)

                else
                    dst[i] = arg
                end
            end
            return concat(dst)
        end

        local function output(...)
            local args = {...}

            return stdout:write(expand_table(args, true))
        end

        ngx.print = output
        ngx.say = function (...)
                local ok, err = output(...)
                if ok then
                    return output("\\\\n")
                end
                return ok, err
            end
        print = ngx.say

        ngx.flush = function (...) return stdout:flush() end
        -- we cannot close stdout here due to a bug in Lua:
        ngx.eof = function (...) return true end
        ngx.exit = os.exit
    ';

    init_worker_by_lua '
        local exit = os.exit
        local stderr = io.stderr

        local function handle_err(err)
            if err then
                err = string.gsub(err, "^init_worker_by_lua:%d+: ", "")
                stderr:write(err, "\\\\n")
            end
            return exit(1)
        end

        local ok, err = pcall(function ()
            if not ngx.config
               or not ngx.config.ngx_lua_version
               or ngx.config.ngx_lua_version < 9011
            then
                error("at least ngx_lua 0.9.12 is required")
            end

$loader
            -- print("calling timer.at...")
            local ok, err = ngx.timer.at(0, function ()
                -- io.stderr:write("timer firing")
                local ok, err = pcall(gen)
                if not ok then
                    return handle_err(err)
                end
                local rc = err
                if rc and type(rc) ~= "number" then
                    return handle_err("bad return value of type " .. type(rc))
                end
                return exit(rc)
            end)
            if not ok then
                return handle_err(err)
            end
            -- print("timer created")
        end)

        if not ok then
            return handle_err(err)
        end
    ';
}
_EOC_

close $out;

my @cmd = ($nginx_path, '-p', "$prefix_dir/", '-c', "conf/nginx.conf");

if ($use_valgrind) {
    my @new = ('valgrind');
    if ($valgrind_opts) {
        $valgrind_opts =~ s/^\s+|\s+$//g;
        push @new, split /\s+/, $valgrind_opts;
    }
    unshift @cmd, @new;
}

my $child_pid;

sub sigint {
    $SIG{INT} = \&sigint;
    if ($child_pid) {
        kill INT => $child_pid;
    }
}
$SIG{INT} = \&sigint;

my $pid = fork();

if (!defined $pid) {
    die "fork() failed: $!\n";
}

if ($pid == 0) {  # child process
    #warn "exec @cmd...";
    exec(@cmd)
        or die "Failed to run command \"@cmd\": $!\n";

} else {
    $child_pid = $pid;
    waitpid($child_pid, 0);
    my $rc = 0;
    if (defined $?) {
        $rc = ($? >> 8);
    }
    exit($rc);
}

sub usage {
    my $rc = shift;
    my $msg = <<_EOC_;
resty [options] [lua-file [args]]

Options:
    -c num              Set maximal connection count (default: 64).
    -e prog             Run the inlined Lua code in "prog".
    --help              Print this help.

    --http-include path Include the specified file in the nginx http configuration block
                        (multiple instances are supported).

    -I dir              Add dir to the search paths for Lua libraries.

    --main-include path Include the specified file in the nginx main configuration block
                        (multiple instances are supported).

    --nginx             Specify the nginx path (this option might be removed in the future).
    -V                  Print version numbers and nginx configurations.
    --valgrind          Use valgrind to run nginx
    --valgrind-opts     Pass extra options to valgrind

For bug reporting instructions, please see:

    <https://openresty.org/en/community.html>

Copyright (C) Yichun Zhang (agentzh). All rights reserved.
_EOC_
    if ($rc == 0) {
        print $msg;
        exit(0);
    }

    warn $msg;
    exit($rc);
}

sub get_bracket_level {
    my %bracket_levels;
    my $bracket_level = 0;
    my $max_level = 0;

    # scan all args and store level of closing brackets
    for my $arg (@_) {
        while ($arg =~ /\](=*)\]/g) {
            my $level = length($1);
            if ($level > $max_level) {
                $max_level = $level;
            }
            $bracket_levels{$level} = 1;
        }
    }

    # if args contain closing bracket
    if (%bracket_levels) {
        # find the shortest form of the long brackets accordingly
        for (my $i = 1; $i < $max_level; $i++) {
            if (!exists $bracket_levels{$i}) {
                $bracket_level = $i;
                last;
            }
        }

        if ($bracket_level == 0) {
            $bracket_level = $max_level + 1;
        }
        return $bracket_level;
    }
    return 1;
}

sub quote_as_lua_str {
    my ($str) = @_;
    my $bracket_level = get_bracket_level($str);
    my $left_bracket = "[" . "=" x $bracket_level . "[";
    my $right_bracket = "]" . "=" x $bracket_level . "]";

    return $left_bracket . $str . $right_bracket;
}

sub gen_lua_code_for_args {
    my ($user_args, $all_args) = @_;

    my $luasrc = "arg = {}\n";

    # args[n] (n = 0)
    $luasrc .= "arg[0] = $quoted_luafile\n";

    # args[n] (n > 0)
    for my $i (0 .. $#user_args) {
        my $index = $i + 1;
        my $quoted_arg = quote_as_lua_str($user_args[$i]);
        $luasrc .= "arg[$index] = $quoted_arg\n";
    }

    # args[n] (n < 0)
    my $left_num = $#all_args - $#user_args;
    for my $i (0 .. $left_num - 2) {
        my $index = 0 - $left_num + $i + 1;
        my $quoted_arg = quote_as_lua_str($all_args[$i]);
        $luasrc .= "arg[$index] = $quoted_arg\n";
    }

    # args[n] (n = the index of resty-cli itself)
    my $index = 0 - $left_num;
    my $quoted_arg = quote_as_lua_str($0);
    $luasrc .= "arg[$index] = $quoted_arg\n";

    #warn $luasrc;
    return $luasrc;
}