/*
 * Copyright (c) 2023, Simon Wanner <simon@skyrising.xyz>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/Utf32View.h>
#include <AK/Utf8View.h>
#include <LibUnicode/Punycode.h>

namespace Unicode::Punycode {

// https://www.rfc-editor.org/rfc/rfc3492.html#section-5
static constexpr u32 BASE = 36;
static constexpr u32 TMIN = 1;
static constexpr u32 TMAX = 26;
static constexpr u32 SKEW = 38;
static constexpr u32 DAMP = 700;
static constexpr u32 INITIAL_BIAS = 72;
static constexpr u32 INITIAL_N = 0x80;
static constexpr u32 DELIMITER = '-';

static Optional<u32> digit_value_of_code_point(u32 code_point)
{
    if (code_point >= 'A' && code_point <= 'Z')
        return code_point - 'A';
    if (code_point >= 'a' && code_point <= 'z')
        return code_point - 'a';
    if (code_point >= '0' && code_point <= '9')
        return code_point - '0' + 26;
    return {};
}

static u32 code_point_value_of_digit(u32 digit)
{
    VERIFY(digit < 36);
    if (digit <= 25)
        return 'a' + digit;
    return '0' + digit - 26;
}

// https://www.rfc-editor.org/rfc/rfc3492.html#section-6.1
static u32 adapt(u32 delta, u32 num_points, bool first_time)
{
    // if firsttime then let delta = delta div damp
    if (first_time)
        delta = delta / DAMP;
    // else let delta = delta div 2
    else
        delta = delta / 2;

    // let delta = delta + (delta div numpoints)
    delta = delta + (delta / num_points);

    // let k = 0
    u32 k = 0;

    // while delta > ((base - tmin) * tmax) div 2 do begin
    while (delta > ((BASE - TMIN) * TMAX) / 2) {
        // let delta = delta div (base - tmin)
        delta = delta / (BASE - TMIN);

        // let k = k + base
        k = k + BASE;
    }

    // return k + (((base - tmin + 1) * delta) div (delta + skew))
    return k + (((BASE - TMIN + 1) * delta) / (delta + SKEW));
}

// https://www.rfc-editor.org/rfc/rfc3492.html#section-6.2
ErrorOr<String> decode(StringView input)
{
    size_t consumed = 0;

    // let n = initial_n
    Checked<size_t> n = INITIAL_N;

    // let i = 0
    Checked<u32> i = 0;

    // let bias = initial_bias
    u32 bias = INITIAL_BIAS;

    // let output = an empty string indexed from 0
    Vector<u32> output;

    // consume all code points before the last delimiter (if there is one)
    //   and copy them to output, fail on any non-basic code point
    Optional<size_t> last_delimiter_index = input.find_last(DELIMITER);
    if (last_delimiter_index.has_value()) {
        for (; consumed < last_delimiter_index.value(); consumed++) {
            if (!is_ascii(input[consumed]))
                return Error::from_string_literal("Unexpected non-basic code point");
            TRY(output.try_append(input[consumed]));
        }

        // if more than zero code points were consumed then consume one more
        //   (which will be the last delimiter)
        if (last_delimiter_index.value() > 0) {
            auto next = input[consumed++];
            VERIFY(next == DELIMITER);
        }
    }

    // while the input is not exhausted do begin
    while (consumed < input.length()) {
        // let oldi = i
        Checked<u32> old_i = i;

        // let w = 1
        Checked<u32> w = 1;

        // for k = base to infinity in steps of base do begin
        for (size_t k = BASE;; k += BASE) {
            // consume a code point, or fail if there was none to consume
            if (consumed >= input.length())
                return Error::from_string_literal("No more code points to consume");
            auto code_point = input[consumed++];

            // let digit = the code point's digit-value, fail if it has none
            auto digit = digit_value_of_code_point(code_point);
            if (!digit.has_value())
                return Error::from_string_literal("Invalid base-36 digit");

            // let i = i + digit * w, fail on overflow
            i = i + Checked(digit.value()) * w;
            if (i.has_overflow())
                return Error::from_string_literal("Numeric overflow");

            // let t = tmin if k <= bias {+ tmin}, or
            //         tmax if k >= bias + tmax, or k - bias otherwise
            u32 t = k <= bias ? TMIN : (k >= bias + TMAX ? TMAX : k - bias);

            // if digit < t then break
            if (digit.value() < t)
                break;

            // let w = w * (base - t), fail on overflow
            w = w * Checked(BASE - t);
            if (w.has_overflow())
                return Error::from_string_literal("Numeric overflow");
        }
        // let bias = adapt(i - oldi, length(output) + 1, test oldi is 0?)
        bias = adapt((i - old_i).value(), output.size() + 1, !old_i);

        // let n = n + i div (length(output) + 1), fail on overflow
        n = n + Checked(static_cast<size_t>(i.value() / static_cast<u32>(output.size() + 1)));
        if (n.has_overflow())
            return Error::from_string_literal("Numeric overflow");

        // let i = i mod (length(output) + 1)
        i = i % Checked(static_cast<u32>(output.size() + 1));

        // {if n is a basic code point then fail}
        // NOTE: The full statement enclosed in braces (checking whether n is a basic code point) can be omitted if initial_n exceeds all basic code points
        //       (which is true for Punycode), because n is never less than initial_n.
        VERIFY(!is_ascii(n.value()));

        // insert n into output at position i
        TRY(output.try_insert(i.value(), n.value()));

        // increment i
        i++;
    }

    StringBuilder builder;
    TRY(builder.try_append(Utf32View(output.data(), output.size())));
    return builder.to_string();
}

static Optional<u32> find_smallest_code_point_greater_than_or_equal(Utf32View code_points, u32 threshold)
{
    Optional<u32> result;
    for (auto code_point : code_points) {
        if (code_point >= threshold && (!result.has_value() || code_point < result.value()))
            result = code_point;
    }
    return result;
}

ErrorOr<String> encode(StringView input)
{
    Vector<u32> code_points;
    for (auto code_point : Utf8View(input))
        TRY(code_points.try_append(code_point));
    return encode(Utf32View(code_points.data(), code_points.size()));
}

// https://www.rfc-editor.org/rfc/rfc3492.html#section-6.3
ErrorOr<String> encode(Utf32View input)
{
    Vector<u32> output;

    // let n = initial_n
    Checked<size_t> n = INITIAL_N;

    // let delta = 0
    Checked<size_t> delta = 0;

    // let bias = initial_bias
    u32 bias = INITIAL_BIAS;

    // let h = b = the number of basic code points in the input
    // copy them to the output in order, followed by a delimiter if b > 0
    size_t b = 0;
    for (auto code_point : input) {
        if (is_ascii(code_point)) {
            TRY(output.try_append(code_point));
            b++;
        }
    }
    auto h = b;
    if (b > 0)
        TRY(output.try_append(DELIMITER));

    // while h < length(input) do begin
    while (h < input.length()) {
        // let m = the minimum {non-basic} code point >= n in the input
        auto m = find_smallest_code_point_greater_than_or_equal(input, n.value());
        VERIFY(m.has_value());

        // let delta = delta + (m - n) * (h + 1), fail on overflow
        delta = delta + (Checked(static_cast<size_t>(m.value())) - n) * Checked(h + 1);
        if (delta.has_overflow())
            return Error::from_string_literal("Numeric overflow");

        // let n = m
        n = m.value();

        // for each code point c in the input (in order) do begin
        for (auto c : input) {
            // if c < n {or c is basic} then increment delta, fail on overflow
            if (c < n.value()) {
                delta++;
                if (delta.has_overflow())
                    return Error::from_string_literal("Numeric overflow");
            }

            // if c == n then begin
            if (c == n.value()) {
                // let q = delta
                auto q = delta.value();

                // for k = base to infinity in steps of base do begin
                for (size_t k = BASE;; k += BASE) {
                    // let t = tmin if k <= bias {+ tmin}, or
                    //         tmax if k >= bias + tmax, or k - bias otherwise
                    u32 t = k <= bias ? TMIN : (k >= bias + TMAX ? TMAX : k - bias);

                    // if q < t then break
                    if (q < t)
                        break;

                    // output the code point for digit t + ((q - t) mod (base - t))
                    auto digit = t + ((q - t) % (BASE - t));
                    TRY(output.try_append(code_point_value_of_digit(digit)));

                    // let q = (q - t) div (base - t)
                    q = (q - t) / (BASE - t);
                }
                // output the code point for digit q
                TRY(output.try_append(code_point_value_of_digit(q)));

                // let bias = adapt(delta, h + 1, test h equals b?)
                bias = adapt(delta.value(), h + 1, h == b);

                // let delta = 0
                delta = 0;

                // increment h
                h++;
            }
        }

        // increment delta and n
        delta++;
        n++;
    }

    StringBuilder builder;
    TRY(builder.try_append(Utf32View(output.data(), output.size())));
    return builder.to_string();
}

}