#include <algorithm>
#include <iostream>
#include <optional>
#include <sstream>

#include "check.h"
#include "error.h"
#include "machine.h"
#include "katom.h"
#include "katom_process.h"
#include "util.h"

namespace {

// One application's argument shape, as written: how many positional parts it
// supplies and which option names it names.  Both are counted at nesting
// depth 0 within the application's span, so a bar or an option name belonging
// to a nested klammer is not miscounted as this one's.
//
// This is the same rule the engine uses at run time, but it has to be stated
// again here rather than reused: argument_split() walks the range flatly,
// which is correct THERE because application is post-order -- by the time a
// klammer is applied its nested spans have already been reduced to text.  At
// check time nothing has been reduced, so the nesting is still present and
// must be tracked.  (The depth-0 rule is the same one cond_separator_bars()
// applies for @cond; see doc/cond_evaluation_order.md.)
// A value as written, and whether it is LITERAL: made of nothing but text
// and whitespace at depth 0, so that what the writer typed is what the
// argument type will see.  A value holding a nested klammer, a *variable*,
// a quoted special or a caret form is computed at run time and cannot be
// judged here -- the static-structure rule: check what is fixed, skip what
// is computed.
struct Argument_value
{
    std::string m_text {};
    bool m_literal { true };
};

struct Application_shape
{
    int m_positional { 0 };
    std::vector<std::string> m_options {};
    std::vector<Argument_value> m_positional_values {};
    std::vector<std::pair<std::string, Argument_value>> m_option_values {};
};

bool is_literal_katom(const Katom& k)
{
    return k.m_type == katom_t::text || k.m_type == katom_t::word
        || k.m_type == katom_t::space || k.m_type == katom_t::newline;
}

bool is_boundary_katom(const Katom& k)
{
    return k.m_type == katom_t::bar || k.m_type == katom_t::option_name;
}

Application_shape application_shape(katom_list::const_iterator begin, katom_list::const_iterator end)
{
    Application_shape shape {};

    auto first = begin;
    while (first != end && first->is_whitespace()) ++first;
    if (first == end) return shape;

    // function_symbol_parts() prepends a synthetic bar when the argument list
    // does not open with an option name, so that content before the first bar
    // counts as a positional part.  Mirror that, or "@f a @" would count zero
    // positional arguments.
    bool in_positional = first->m_type != katom_t::option_name;
    if (in_positional) shape.m_positional = 1;

    // The value being accumulated: a positional part, or the option named
    // last.  Closed at each depth-0 boundary and at the end.
    Argument_value current {};
    bool current_is_option = !in_positional;
    std::string current_option {};
    auto close = [&]() {
        current.m_text = trim(current.m_text);
        if (current_is_option) {
            shape.m_option_values.emplace_back(current_option, current);
        } else {
            shape.m_positional_values.push_back(current);
        }
        current = Argument_value {};
    };

    int depth = 0;
    for (auto k = first; k != end; ++k) {
        if (depth == 0 && is_boundary_katom(*k)) {
            close();
            if (k->m_type == katom_t::bar) {
                ++shape.m_positional;
                current_is_option = false;
            } else {
                shape.m_options.push_back(k->m_text.substr(1));
                current_is_option = true;
                current_option = k->m_text.substr(1);
            }
        } else {
            if (depth > 0 || !is_literal_katom(*k)) current.m_literal = false;
            current.m_text += k->m_text;
        }
        if (level_increase(*k)) {
            ++depth;
        } else if (level_decrease(*k)) {
            --depth;
        }
    }
    close();
    return shape;
}

// The span of the application opening at `begin`, as [begin, end): end is one
// past the matching close.  Empty when the span is unclosed -- which the
// engine reports separately, so the checker just stops descending.
//
// The result must be an optional rather than "list_end means unclosed": a
// span that closes on the very last katom of the list -- a klammer body that
// is nothing but one application, "@@u : @nosuch x @ @@" -- ends exactly AT
// list_end while being perfectly well formed, and conflating the two made the
// checker silently skip every such body.
std::optional<katom_list::const_iterator> span_end(
    katom_list::const_iterator begin, katom_list::const_iterator list_end)
{
    int depth = 0;
    for (auto k = begin; k != list_end; ++k) {
        if (level_increase(*k)) {
            ++depth;
        } else if (level_decrease(*k)) {
            if (--depth == 0) return k + 1;
        }
    }
    return {};
}

bool skip_katom(const Katom& k)
{
    return k.m_type == katom_t::replaced
        || k.m_type == katom_t::ignored
        || k.m_type == katom_t::literal;
}

// Argument spans of the primitives whose contents are not Klammertext: @eval
// receives code, @read a filename.  @cond is NOT in this set -- its branches
// are Klammertext, and checking the branch that is not selected is the main
// thing the checker is for.
bool opens_uncheckable_span(const Katom& k)
{
    return k.m_type == katom_t::eval_begin || k.m_type == katom_t::read_begin;
}

class Checker
{
public:
    Checker(Machine& machine, std::vector<Diagnostic>& diagnostics)
        : m_machine(machine)
        , m_diagnostics(diagnostics)
        {}

    void check_list(const katom_list& katoms, const std::string& target,
                    const std::string& context);

private:
    void check_application(
        const std::string& name, const Klammer& klammer,
        katom_list::const_iterator begin, katom_list::const_iterator end,
        const std::string& target, const std::string& context);

    void error(const std::string& message, const std::string& context, const Locator& loc)
    {
        m_diagnostics.emplace_back("error", message, context, loc);
    }

    Machine& m_machine;
    std::vector<Diagnostic>& m_diagnostics;
};

void Checker::check_application(
    const std::string& name, const Klammer& klammer,
    katom_list::const_iterator begin, katom_list::const_iterator end,
    const std::string& target, const std::string& context)
{
    // Arity is a property of the klammer's rationalized parameter set, which
    // is shared by all of its target definitions, so it is checked once here
    // rather than per target.
    const Parameter_set& parameters = klammer.m_parameters;
    Application_shape shape = application_shape(begin + 1, end - 1);

    auto required = static_cast<int>(parameters.m_positional.size());
    bool has_rest = !parameters.m_rest.empty();

    if (shape.m_positional < required) {
        std::stringstream ss {};
        ss << "@" << name << " needs " << required << " positional "
           << plural("argument", required) << " but is given " << shape.m_positional
           << ".  Positional arguments are separated by \"|\".";
        error(ss.str(), context, begin->m_loc);
    } else if (shape.m_positional > required && !has_rest) {
        std::stringstream ss {};
        ss << "@" << name << " takes " << required << " positional "
           << plural("argument", required) << " but is given " << shape.m_positional << ".";
        error(ss.str(), context, begin->m_loc);
    }

    // Argument VALUES against their argument types -- the check that finds
    // a height unit on a width, a word where a number is wanted -- for every
    // value that is literal (see Argument_value).  The same test the engine
    // applies when the klammer is applied, Parameter_set::validate(), so the
    // two can never disagree; its exception becomes a diagnostic here.
    auto validate = [&](const Parameter& parameter, const Argument_value& value) {
        if (!value.m_literal || value.m_text.empty()) return;
        try {
            Parameter_set::validate(parameter, value.m_text, begin->m_loc);
        } catch (const Argument_error& e) {
            error("@" + name + ": " + collapse_whitespace(e.m_desc), context, begin->m_loc);
        }
    };
    for (size_t i = 0; i < shape.m_positional_values.size() && i < parameters.m_positional.size(); i++) {
        validate(parameters.m_positional[i], shape.m_positional_values[i]);
    }
    for (const auto& [option, value] : shape.m_option_values) {
        if (const Parameter* parameter = parameters.find(option)) {
            validate(*parameter, value);
        }
    }

    std::vector<std::string> seen {};
    for (const auto& option : shape.m_options) {
        if (std::ranges::count(parameters.m_optional_names, option) == 0) {
            std::stringstream ss {};
            ss << "@" << name << " has no optional argument \":" << option << "\".";
            if (!parameters.m_optional_names.empty()) {
                ss << "  It accepts: :" << join(parameters.m_optional_names, " :") << ".";
            }
            error(ss.str(), context, begin->m_loc);
        } else if (std::ranges::count(seen, option) > 0) {
            error("@" + name + " is given \":" + option + "\" more than once.",
                  context, begin->m_loc);
        }
        seen.push_back(option);
    }

    // Target coverage.  A klammer may be declared (.k) and defined for some
    // targets but not the one being built; run time only discovers this if the
    // application is actually reached.
    //
    // Not checked under the general target: a general body is not applied
    // under "*", it is copied to every target that lacks its own definition
    // and applied under whichever of those is in force (copy_general_klammer_
    // to_undefined() in klammer.cpp).  So an application inside it resolves
    // against a real target, and the per-target passes are where coverage is
    // decided.  Checking it here reported @b -- defined for html/tex/pdf/txt
    // but not for "*" -- as missing from a general body that in fact works.
    if (target != Target_registry::general_name && klammer.m_defloc.count(target) == 0) {
        std::stringstream ss {};
        ss << "@" << name << " is not defined for the target \"" << target << "\".";
        strings_t targets = klammer.get_target_names();
        if (!targets.empty()) {
            ss << "  It is defined for: " << join(targets, ", ") << ".";
        }
        error(ss.str(), context, begin->m_loc);
    }
}

void Checker::check_list(
    const katom_list& katoms, const std::string& target, const std::string& context)
{
    for (auto k = katoms.begin(); k != katoms.end(); ++k) {
        if (k->m_type == katom_t::ignore_rest) break;
        if (skip_katom(*k)) continue;

        // Code and filenames, not applications: skip the whole span.
        if (opens_uncheckable_span(*k)) {
            auto skip_to = span_end(k, katoms.end());
            if (!skip_to) return;
            k = *skip_to - 1;
            continue;
        }
        if (k->m_type != katom_t::apply_begin) continue;

        auto closed = span_end(k, katoms.end());
        if (!closed) return;               // unclosed; the engine reports it
        auto end = *closed;

        std::string name = trim_char(k->m_text, '@');
        auto found = m_machine.m_klammers.m_klammers.find(name);
        if (found == m_machine.m_klammers.m_klammers.end()) {
            error("The klammer @" + name + " is not defined.", context, k->m_loc);
            continue;
        }
        check_application(name, found->second, k, end, target, context);

        // A literal parameter's content is raw text -- a "@" inside it is not
        // an application -- so do not descend into it.
        if (found->second.has_literal_param()) {
            k = end - 1;
        }
    }
}

} // namespace

std::vector<Diagnostic> check_machine(Machine& machine, const katom_list& document,
                                      const std::string& target)
{
    std::vector<Diagnostic> diagnostics {};
    Checker checker(machine, diagnostics);

    // With no target named, check every target the machine defines, plus the
    // general one -- a klammer defined without a target suffix has its body
    // filed under the general name, and with no klammer set loaded that is the
    // only target there is.
    strings_t targets {};
    if (target == Target_registry::general_name) {
        targets = machine.m_targets.user_defined();
        targets.push_back(Target_registry::general_name);
    } else {
        targets.push_back(target);
    }

    for (const auto& t : targets) {
        checker.check_list(document, t, "document");
        for (const auto& [name, klammer] : machine.m_klammers.m_klammers) {
            auto body = klammer.m_body.find(t);
            if (body == klammer.m_body.end()) continue;
            checker.check_list(body->second, t, "body of @" + name);
        }
    }

    // The same text is checked once per target, so a fault that does not
    // depend on the target -- an undefined name, a wrong argument count --
    // is found once per target and must be reported once.  Target coverage
    // names its target in the message, so those stay distinct.  Hence the
    // context deliberately does NOT carry the target: it is what makes the
    // target-independent duplicates compare equal.
    std::vector<Diagnostic> unique {};
    for (const auto& d : diagnostics) {
        bool seen = std::any_of(
            unique.begin(), unique.end(), [&d](const Diagnostic& u) {
                return u.m_severity == d.m_severity && u.m_message == d.m_message
                    && u.m_context == d.m_context && u.m_loc.str() == d.m_loc.str(); });
        if (!seen) unique.push_back(d);
    }
    return unique;
}

int report_diagnostics(const std::vector<Diagnostic>& diagnostics, std::ostream& os)
{
    int errors = 0;
    for (const auto& d : diagnostics) {
        if (d.m_severity == "error") ++errors;
        os << d.m_severity << ": " << d.m_message << "\n";
        if (!d.m_context.empty()) {
            os << "  in " << d.m_context << "\n";
        }
        if (!d.m_loc.m_filename.empty()) {
            os << "  " << d.m_loc.desc() << "\n";
        }
        os << "\n";
    }
    os << diagnostics.size() << " " << plural("diagnostic", diagnostics.size())
       << ", " << errors << " " << plural("error", errors) << "\n";
    return errors;
}