-- Load libraries
local exit = require "sm.utils.exit"
local global_config = require "sm.config"
local redis = require "sm.utils.redis"
local config_fetcher = require "sm.utils.config_fetcher"
local random = require "sm.utils.random"
local cjson = require "cjson"
local resolver = require "resty.dns.resolver"
local hls = require "sm.utils.hls"

-- Reference native functions
local os_getenv = os.getenv
local NODE_TYPE = os_getenv("NODE_TYPE")
local ngx_unescape_uri = ngx.unescape_uri
local VERSION = _VERSION

-- Random seed
random.seed()

local M = {}

-- 连接Redis
local function connect_redis()
    local remote_addr = ngx.var.remote_addr
    local request_id = ngx.var.request_id

    local ok, err = redis.connect(global_config["redis"])
    if not ok then
        ngx.log(ngx.ERR, "Failed to connect to Redis: " .. err)
        return nil, exit.error(remote_addr, request_id)
    end
    return true, nil
end

-- 判断是否为回源节点
local function is_origin_node()
    return NODE_TYPE == "origin"
end

-- 设置后端协议
local function set_backend_protocol(hostname)
    if hostname["origin_scheme"] then
        ngx.var.backend_protocol = hostname["origin_scheme"] .. "://"
    else
        ngx.var.backend_protocol = "http://"
    end
    ngx.log(ngx.INFO, "backend_protocol: " .. ngx.var.backend_protocol)
end

-- 创建DNS解析器
local function create_dns_resolver()
    local remote_addr = ngx.var.remote_addr
    local request_id = ngx.var.request_id

    local r, err = resolver:new{
        nameservers = {"8.8.8.8", {"8.8.4.4", 53}},
        retrans = 5, -- 5 retransmissions on receive timeout
        timeout = 2000, -- 2 sec
        no_random = true -- always start with first nameserver
    }

    if not r then
        ngx.log(ngx.ERR, "Failed to create resolver: " .. err)
        return nil, exit.error(remote_addr, request_id)
    end

    return r, nil
end

-- 获取源站配置
local function get_sources_config(zone_id)
    local remote_addr = ngx.var.remote_addr
    local request_id = ngx.var.request_id

    local json_sources, hit_level, err = config_fetcher.getDomainConf(redis, zone_id .. ":sources")
    if err then
        ngx.log(ngx.ERR, "Failed to get zone config: " .. zone_id .. " ,error: " .. cjson.encode(err))
        return nil, exit.error(remote_addr, request_id)
    end

    local sources = cjson.decode(json_sources)
    ngx.log(ngx.INFO, "sources: " .. json_sources)
    return sources, nil
end

-- 获取回源节点列表（非回源节点使用）
local function get_origin_nodes_config(zone_id)
    local remote_addr = ngx.var.remote_addr
    local request_id = ngx.var.request_id

    -- 查询回源节点列表
    local json_origin_nodes, hit_level, err = config_fetcher.getDomainConf(redis, zone_id .. ":origin_node_list")
    if err then
        ngx.log(ngx.ERR, "Failed to get origin nodes config: " .. zone_id .. " ,error: " .. cjson.encode(err))
        return nil, exit.error(remote_addr, request_id)
    end

    if not json_origin_nodes then
        ngx.log(ngx.ERR, "No origin nodes found for zone: " .. zone_id)
        return nil, exit.error(remote_addr, request_id)
    end

    local origin_nodes = cjson.decode(json_origin_nodes)
    ngx.log(ngx.INFO, "origin_nodes: " .. json_origin_nodes)
    return origin_nodes, nil
end

-- 处理DNS解析
local function resolve_addresses(sources, dns_resolver, hostname)
    for _, item in ipairs(sources) do
        local ip = item.address

        if item.resolve then
            -- TODO: 按照ttl配置缓存
            local ips, err, tries = dns_resolver:query(item.address, nil, {})
            if err then
                ngx.log(ngx.ERR, "[Balancer] Failed to resolve address: " .. item.address .. " " .. err)
            elseif not ips or #ips == 0 then
                ngx.log(ngx.ERR, "[Balancer] No IPs resolved for address: " .. item.address)
            else
                ngx.log(ngx.INFO, "ips: " .. cjson.encode(ips))
                -- 查找第一个有address字段的记录
                local found_ip = nil
                for _, v in ipairs(ips) do
                    if v.address then
                        found_ip = v.address
                        break
                    end
                end
                if found_ip then
                    ngx.log(ngx.INFO, "[Balancer] Resolved address: " .. item.address .. " to " .. found_ip)
                    ip = found_ip
                else
                    ngx.log(ngx.ERR, "[Balancer] No valid address found in DNS response for: " .. item.address)
                end
            end
            if hostname["sni"] then
                item.sni = hostname["sni"]
            end
        end

        item.ip = ip
    end
end

-- 设置Proxy Host
local function set_proxy_host(hostname)
    -- 根据节点类型设置Proxy Host
    if is_origin_node() or (NODE_TYPE == nil) then
        local proxy_host = hostname["proxy_host"]
        if proxy_host and proxy_host ~= "$host" then
            ngx.log(ngx.INFO, "proxy_host: " .. proxy_host)
            ngx.var.proxy_host = proxy_host
        end
    end
end

-- 处理回源节点逻辑
local function handle_origin_node(zone_id, hostname)
    ngx.log(ngx.INFO, "Processing as origin node")

    -- 获取源站配置
    local sources, err = get_sources_config(zone_id)
    if not sources then
        return err
    end

    -- 创建DNS解析器
    local dns_resolver, err = create_dns_resolver()
    if not dns_resolver then
        return err
    end

    -- 处理DNS解析
    resolve_addresses(sources, dns_resolver, hostname)

    -- 设置源站列表
    ngx.var.sources = cjson.encode(sources)

    return nil
end

-- 处理非回源节点逻辑
local function handle_non_origin_node(zone_id, hostname)
    ngx.log(ngx.INFO, "Processing as non-origin node")

    -- 获取回源节点列表
    local origin_nodes, err = get_origin_nodes_config(hostname["region_id"])
    if not origin_nodes then
        return err
    end

    -- 设置回源节点列表
    ngx.var.sources = cjson.encode(origin_nodes)

    return nil
end

-- 获取鉴权配置
-- Redis中的配置格式应为 zone_id:auth
-- 配置示例：
-- {
--   "enable": 1,           -- 是否开启 1-开启 2-关闭
--   "auth_type": 1,         -- 鉴权类型1方式A
--   "primary_key": "key123", -- 主key
--   "backup_key": "backup", -- 备key
--   "auth_m3u8": true,      -- 鉴权继承配置
--   "auth_delta": 1800,     -- 有效时间，默认1800秒
--   "match": "%.jpg$,%.js$,%.ts$", -- 鉴权范围，文件后缀
--   "algorithm": "sha256"  -- 加密算法 md5/sha256
-- }
local function get_auth_config(zone_id)
    local json_auth_config, hit_level, err = config_fetcher.getDomainConf(redis, zone_id .. ":auth_conf")
    if not json_auth_config then
        ngx.log(ngx.INFO, "No auth config found for zone: " .. zone_id)
        return nil
    end
    
    local auth_config = cjson.decode(json_auth_config)
    ngx.log(ngx.INFO, "auth_config: " .. json_auth_config)
    return auth_config
end

-- 执行鉴权检查（仅非回源节点）
local function perform_auth_check(zone_id)
    if is_origin_node() then
        ngx.log(ngx.INFO, "Origin node detected, skipping authentication check")
        return nil
    end
    
    ngx.log(ngx.INFO, "Non-origin node, performing authentication check")
    
    -- 获取鉴权配置
    local auth_config = get_auth_config(zone_id)
    if not auth_config then
        ngx.log(ngx.INFO, "No auth config, skipping authentication")
        return nil
    end
    
    local result = auth_check(auth_config)
    return result
end

-- 完成访问处理
local function finalize_access(zone_id, hostname)
    -- 自定义回源请求头
    M.custom_headers(zone_id)

    -- 设置Proxy Host
    set_proxy_host(hostname)

    -- 关闭Redis连接
    redis.close()
end

function M.access()
    local host = ngx.var.host
    local remote_addr = ngx.var.remote_addr
    local request_id = ngx.var.request_id
    local zone_id = ngx.var.zone_id

    ngx.log(ngx.INFO, "hostname: " .. host .. ", node_type: " .. (NODE_TYPE or "undefined"))
    local hostname = cjson.decode(ngx.var.site)

    -- 设置后端协议
    set_backend_protocol(hostname)

    -- 连接Redis
    local ok, err = connect_redis()
    if not ok then
        return err
    end

        -- 执行鉴权检查（仅非回源节点）
    local auth_result = perform_auth_check(zone_id)
    if auth_result then
        return auth_result
    end

    -- 根据节点类型处理不同逻辑
    local result
    if is_origin_node() then
        result = handle_origin_node(zone_id, hostname)
    elseif NODE_TYPE == "edge" then
        result = handle_non_origin_node(zone_id, hostname)
    else
        result = handle_origin_node(zone_id, hostname)
    end
    
    if result then
        return result
    end

    -- 完成访问处理
    finalize_access(zone_id, hostname)
end

-- 自定义回源请求头
function M.custom_headers(zone_id)
    -- Get zone config
    local json_custom_headers, hit_level, err =
        config_fetcher.getDomainConf(redis, zone_id .. ":origin_request_headers")
    if json_custom_headers then
        local custom_headers = cjson.decode(json_custom_headers)
        ngx.log(ngx.INFO, "custom_headers: " .. json_custom_headers)
        if custom_headers then
            for _, header in ipairs(custom_headers or {}) do
                ngx.req.set_header(header.key, header.value)
            end
        end
    end

    -- Get zone config
    local json_origin_response_headers, hit_level, err = config_fetcher.getDomainConf(redis, zone_id ..
        ":origin_response_headers")
    if json_origin_response_headers then
        local origin_response_headers = cjson.decode(json_origin_response_headers)
        if origin_response_headers then
            ngx.var.origin_response_headers = origin_response_headers
        end
    end
end

-- 鉴权判断
function auth_check(auth_config)
    -- 检查是否启用鉴权 1-开启 2-关闭
    if auth_config.enable ~= 1 then
        ngx.log(ngx.INFO, "Authentication disabled")
        return nil
    end
    
    -- 根据鉴权类型执行相应的鉴权逻辑
    local auth_type = auth_config.auth_type
    local result
    if auth_type == 1 then
        -- 鉴权类型A
        result = auth_a(auth_config)
    elseif auth_type == 2 then
        -- 鉴权类型B
        result = auth_b(auth_config)
    elseif auth_type == 3 then
        -- 鉴权类型C
        result = auth_c(auth_config)
    elseif auth_type == 4 then
        -- 鉴权类型F
        result = auth_f(auth_config)
    else
        ngx.log(ngx.ERR, "Unknown auth type: " .. auth_type)
        local remote_addr = ngx.var.remote_addr
        local request_id = ngx.var.request_id
        return exit.rule_block(remote_addr, request_id)
    end
    
    return result
end

-- 解析文件后缀配置字符串
-- 数据已经是处理好的lua正则表达式，只需按逗号拆分
local function parse_suffixes(match_config)
    if not match_config or match_config == "" then
        return {}
    end
    
    local suffixes = {}
    -- 按逗号分割正则表达式
    for pattern in match_config:gmatch("[^,]+") do
        pattern = pattern:gsub("^%s*(.-)%s*$", "%1") -- 去除前后空格
        if pattern ~= "" then
            table.insert(suffixes, pattern)
            ngx.log(ngx.DEBUG, "Added suffix pattern: " .. pattern)
        end
    end
    
    ngx.log(ngx.INFO, "Parsed " .. #suffixes .. " suffix patterns from: " .. match_config)
    return suffixes
end

-- 鉴权类型A
function auth_a(auth_config)
    local remote_addr = ngx.var.remote_addr
    local request_id = ngx.var.request_id
    
    -- 从配置中获取参数
    local auth_m3u8 = auth_config.auth_m3u8 or false -- 是否对m3u8文件继承鉴权参数
    local auth_sign = auth_config.algorithm or "md5" -- 加密方法
    local primary_key = auth_config.primary_key or ""
    local backup_key = auth_config.backup_key or ""
    local valid_duration = auth_config.auth_delta or 1800 -- 鉴权URL有效时长，单位秒
    
    ngx.log(ngx.INFO, string.format("Auth config - m3u8: %s, algorithm: %s, duration: %d", 
        tostring(auth_m3u8), auth_sign, valid_duration))
    
    -- 解析文件后缀配置
    local suffixes = parse_suffixes(auth_config.match)
    local contains_m3u8 = false
    
    -- 判断是否需要鉴权,如果列表为空，则默认鉴权
    if #suffixes > 0 then
        -- 检查 URI 是否匹配指定的后缀
        local matched = false
        for _, pattern in ipairs(suffixes) do
            if ngx.var.uri:match(pattern) then
                matched = true
            end
            if pattern == "%.m3u8$" then
                contains_m3u8 = true
            end
        end
        if not matched then
            ngx.log(ngx.INFO, "URI does not match auth suffixes, skipping auth")
            return nil
        end
    else
        contains_m3u8 = true
    end
    -- 获取请求参数
    local uri = ngx.var.uri
    local args = ngx.req.get_uri_args()
    local auth_key = args["auth_key"]
    -- 检查是否包含鉴权参数
    if not auth_key then
        ngx.status = ngx.HTTP_FORBIDDEN
        ngx.log(ngx.DEBUG, "Missing auth_key")
        return exit.rule_block(remote_addr, request_id)
    end
    -- 解析鉴权参数
    local timestamp, rand, uid, hash = auth_key:match("(%d+)-(%w+)-(%d+)-(%w+)")
    -- 检查匹配结果是否为 nil
    if not timestamp or not rand or not uid or not hash then
        ngx.status = ngx.HTTP_FORBIDDEN
        ngx.log(ngx.DEBUG, "Invalid auth_key format")
        return exit.rule_block(remote_addr, request_id)
    end
    -- 确保 uid 变量已定义并初始化
    uid = uid or "0" -- 如果 uid 为 nil，则赋值为 "0"
    -- 检查时间戳是否有效
    local current_time = ngx.time()
    if current_time > tonumber(timestamp) + valid_duration then
        ngx.status = ngx.HTTP_FORBIDDEN
        ngx.log(ngx.DEBUG, "URL expired")
        return exit.rule_block(remote_addr, request_id)
    end
    -- 根据配置选择加密算法
    local sign_func = function(uri_str, timestamp_str, rand_str, uid_str, key_str)
        local sign_string = string.format("%s-%s-%s-%s-%s", uri_str, timestamp_str, rand_str, uid_str, key_str)
        
        if auth_sign == "md5" then
            return hls.md5(sign_string)
        elseif auth_sign == "sha256" then
            return hls.sha256(sign_string)
        else
            -- 默认使用md5
            return hls.md5(sign_string)
        end
    end
    -- 使用主秘钥进行鉴权
    local function authenticate(key)
        if not key or key == "" then
            return false
        end
        local calculated_hash = sign_func(uri, timestamp, rand, uid, key)
        ngx.log(ngx.DEBUG, "Calculated hash: " .. calculated_hash .. ", Expected hash: " .. hash)
        return calculated_hash == hash
    end
    
    -- 检查主备key是否都为空
    if (not primary_key or primary_key == "") and (not backup_key or backup_key == "") then
        ngx.status = ngx.HTTP_FORBIDDEN
        ngx.log(ngx.ERR, "Both primary and backup keys are empty")
        return exit.rule_block(remote_addr, request_id)
    end
    
    -- 主备都没有通过
    local primary_auth = authenticate(primary_key)
    local backup_auth = authenticate(backup_key)
    
    if not primary_auth and not backup_auth then
        ngx.status = ngx.HTTP_FORBIDDEN
        ngx.log(ngx.DEBUG, "Authentication failed for both primary and backup keys")
        return exit.rule_block(remote_addr, request_id)
    end
    
    ngx.log(ngx.INFO, "Authentication successful")
    -- 请求通过，如果配置需要为m3u8文件继承鉴权参数
    ngx.ctx.auth_m3u8 = auth_m3u8
    if auth_m3u8 and contains_m3u8 then
        ngx.ctx.auth_timestamp = timestamp
        ngx.ctx.auth_rand = rand
        ngx.ctx.auth_uid = uid
        if primary_key ~= "" then
            ngx.ctx.auth_key = primary_key
        else
            ngx.ctx.auth_key = backup_key
        end
        ngx.ctx.auth_sign = auth_sign
        ngx.log(ngx.INFO, "Auth parameters inherited for m3u8/ts files")
    end
    
    -- 鉴权成功，返回nil表示继续处理
    return nil
end

return M