#include <set>
#include <utility>
#include "machine.h"
#include "error.h"
#include "show.h"
#include "util.h"
#include "file.h"
#include "log.h"
#include "eval.h"

Machine::Machine()
    : m_argtypes(Argtype_set())
    , m_state(State())
    , m_targets(Target_set())
    , m_klammers(Klammer_set())
{
    (void)K::log(3);
    m_state.add_environment_frame();
    /*
    if (sks) {
        fs::path sks_filename(klammertext_filename("sks/sks.k"));
        // msg() << "SKS filename: " << sks_filename << "\n";
        read(sks_filename);
    }
    */
}

void Machine::process_eval_katoms(katom_list& katoms)
{
    (void)K::log(3);
    if (std::find_if(katoms.begin(), katoms.end(), begin_eval) != katoms.end()) {
        for (const auto& [op, cl] : find_spans(katoms, level_increase, level_decrease, true, "eval")) {
            auto [begin, end] = find_span_katoms(katoms, op, cl);
            if (begin_eval(*begin)) {
                Eval E(*this, begin->m_loc);
                katom_list eval_katoms = E.eval(begin, end);
                std::for_each(begin, end, mark_as_replaced);
                katoms.insert(end, eval_katoms.begin(), eval_katoms.end());
            }
        }
    }
}

// Collect the bars that are direct argument separators of a @cond span:
// the bar katoms at nesting depth 0 within the span.  A bar that lies
// inside a nested span — for example the "|" in an inner @frac a | b @, or
// in a nested @eval/@read/@cond — has positive depth and is excluded.
//
// This makes @cond's argument delimitation a property of the span tree
// (the operad's arity: each operator owns the bars at its own level) rather
// than of the flat katom range.  Counting every bar in the range, as the
// original check did, conflated the arities of nested operators and rejected
// well-formed input such as
//   @cond *bool* | @frac 1 | 2 @ | @frac 2 | 1 @ @
// because the inner @frac bars were miscounted as @cond separators.
//
// begin is the cond_begin katom; end is one past the closing apply_end, so
// *(end - 1) is the apply_end.  Bars are returned in source order.
std::vector<katom_iter> cond_separator_bars(katom_iter begin, katom_iter end)
{
    std::vector<katom_iter> bars {};
    int depth = 0;
    for (auto it = begin + 1; it != end - 1; ++it) {
        if (is_bar(*it) && depth == 0) {
            bars.push_back(it);
        } else if (level_increase(*it)) {
            ++depth;
        } else if (level_decrease(*it)) {
            --depth;
        }
    }
    return bars;
}

void check_bar_count(katom_iter begin, std::size_t count)
{
    if (count != 1 && count != 2) {
        std::stringstream ss {};
        ss << "Incorrectly formatted @cond klammer. There should only be one or two bar characters:\n"
           << "  @cond <predicate> | <result-if-true @\nor:\n"
           << "  @cond <predicate> | <result-if-true> | <result-if-false> @";
        throw Argument_error(ss.str(), begin->m_loc, false);
    }
}

bool is_true(const std::string& s)
{
    return s == "True" || s == "true" || s == "1";
}

void Machine::process_cond_katoms(katom_list& katoms)
{
    if (std::find_if(katoms.begin(), katoms.end(), begin_cond) != katoms.end()) {
        (void)K::log(3);
        //for (auto [op, cl] : find_spans(katoms, begin_cond, end_apply, true, "cond")) {
        for (const auto& [op, cl] : find_spans(katoms, level_increase, level_decrease, true, "cond")) {
            auto [begin, end] = find_span_katoms(katoms, op, cl);
            // msg() << "find_spans: " << std::pair(begin, end) << "\n";
            if (begin_cond(*begin)) {
                // Delimit @cond's arguments by the bars at depth 0 within the
                // span, so that bars belonging to nested klammers are not
                // mistaken for @cond's own separators (see cond_separator_bars).
                std::vector<katom_iter> bars = cond_separator_bars(begin, end);
                check_bar_count(begin, bars.size());
                auto bar_1 = bars[0];
                std::string predicate = to_string(begin + 1, bar_1, true);
                katom_list true_clause {};
                katom_list false_clause {};
                if (bars.size() == 2) {
                    auto bar_2 = bars[1];
                    true_clause = katom_list(bar_1 + 1, bar_2);
                    false_clause = katom_list(bar_2 + 1, end - 1);
                } else {
                    true_clause = katom_list(bar_1 + 1, end - 1);
                }
                // Splice only the selected branch.  Its nested klammers remain
                // unreduced here and are reduced by the outer fixed-point apply
                // loop; the unselected branch is discarded without evaluation
                // (@cond is a non-strict special form).
                katom_list result = is_true(predicate)
                    ? trim_whitespace(true_clause) : trim_whitespace(false_clause);
                std::for_each(begin, end, mark_as_replaced);
                katoms.insert(end, result.begin(), result.end());
            }
        }
    }
}


void Machine::expand_constant_klammers(katom_list& katoms, const Katom& op, const Katom& cl)
{
    auto [begin, end] = find_span_katoms(katoms, op, cl);
    restore_initial_type(begin + 1, end - 1);
    if (std::find_if(begin + 1, end - 1, begin_klammer_apply) == end - 1) return;
    for (const auto& [app_op, app_cl] : find_spans(begin + 1, end - 1, begin_apply, end_apply, false, "def-time")) {
        auto [app_begin, app_end] = find_span_katoms(katoms, app_op, app_cl);
        if (app_begin->m_type == katom_t::apply_begin) {
            std::string name = trim_char(app_begin->m_text, '@');
            const auto* body = m_klammers.constant_body(name);
            if (body) {
                // Set both m_type and m_initial_type so that
                // restore_initial_type() in add() won't resurrect them
                for (auto it = app_begin; it != app_end; ++it) {
                    it->m_type = katom_t::replaced;
                    it->m_initial_type = katom_t::replaced;
                }
                katoms.insert(app_end, body->begin(), body->end());
            }
        }
    }
}

//katom_list
void Machine::mark_literal_klammer_content(katom_list& katoms)
{
    (void)K::log(4);
    // Collect names of klammers that have a literal parameter
    std::set<std::string> literal_names {};
    for (const auto& [name, klammer] : m_klammers.m_klammers) {
        if (klammer.has_literal_param())
            literal_names.insert(name);
    }
    if (literal_names.empty()) return;

    // Scan for matching @name ... name@ spans.
    // Stop at ## (ignore-rest) since everything after it will be removed.
    for (auto k = katoms.begin(); k != katoms.end(); ++k) {
        if (k->m_type == katom_t::ignore_rest) break;
        if (k->m_type != katom_t::apply_begin) continue;
        std::string name = trim_char(k->m_text, '@');
        if (literal_names.count(name) == 0) continue;
        (void)K::log(2, "Literal klammer: " + name);

        // Find the matching named closing delimiter
        std::string close_text = name + "@";
        auto close = k + 1;
        int depth = 1;
        while (close != katoms.end()) {
            if (close->m_type == katom_t::apply_begin &&
                trim_char(close->m_text, '@') == name)
                depth++;
            else if (close->m_type == katom_t::apply_end &&
                     trim_char(close->m_text, '@') == name)
                depth--;
            if (depth == 0) break;
            ++close;
        }
        if (close == katoms.end()) {
            throw Parsing_error(
                "Klammer " + q_(name) + " has a literal parameter and must be closed with "
                + q_(close_text),
                k->m_loc);
        }

        // Count positional parameters before the literal one.
        // The literal parameter is always last. Bars separate
        // the preceding positional arguments and the literal content.
        const auto& klammer = m_klammers.m_klammers[name];
        int bars_before_literal = 0;
        for (const auto& p : klammer.m_parameters.m_positional) {
            if (p.m_argtype.m_name == "literal") break;
            bars_before_literal++;
        }

        // Find where literal content starts.
        // Skip bars_before_literal bars (separating preceding positional args).
        // If options are present, skip past the bar after them.
        // Options are identified by :name katoms before any bar.
        auto literal_start = k + 1;
        bool has_options = false;
        for (auto j = k + 1; j < close; ++j) {
            if (j->m_type == katom_t::option_name) {
                has_options = true;
            }
            if (j->m_type == katom_t::bar) {
                if (bars_before_literal > 0) {
                    bars_before_literal--;
                    literal_start = j + 1;
                } else if (has_options) {
                    // This bar separates options from literal content
                    literal_start = j + 1;
                    break;
                } else {
                    // No preceding args, no options: bar is part of literal
                    break;
                }
            }
        }
        std::for_each(literal_start, close, mark_as_literal);

        k = close;  // Skip past this span
    }
}

void Machine::process_katoms(
    katom_list& katoms, const std::string& source,
    bool nonascii, bool literal, bool ignore, bool whitespace, bool klammers,
    bool eval, bool cond, bool read)
{
    mark_literal_klammer_content(katoms);
    if (literal) mark_literal_katoms(katoms);
    hide_special_katoms(katoms);
    if (nonascii) encode_nonascii_characters(katoms);
    if (ignore) mark_ignored_katoms(katoms);
    if (whitespace) process_whitespace_modifiers(katoms);
    if (klammers) process_klammer_katoms(katoms);
    if (eval) process_eval_katoms(katoms);
    if (cond) process_cond_katoms(katoms);
    if (read) expand_read_katoms(
        katoms, source,
        nonascii, literal, ignore, whitespace, klammers, eval, cond, read);
    warn_unparsed_katoms(katoms, m_warn_unparsed);
//    return katoms;
}

katom_list Machine::process(
    std::string text, const std::string& source,
    bool nonascii, bool literal, bool ignore, bool whitespace, bool klammers,
    bool eval, bool cond, bool read)
{
    katom_list katoms = katomize(line_split(text), source);
    //katoms =
    process_katoms(
        katoms, source,
        nonascii, literal, ignore, whitespace, klammers, eval, cond, read);    
    return katoms;
}

void Machine::read(const fs::path& pathname)
{
    (void)K::log(3, pathname.string());
    m_state.open_frame("Machine state: " + pathname.string());
    std::string text = m_state.subst(trim_right(string_from_file(pathname)));
    katom_list katoms = process(text, pathname);
    m_katoms.insert(m_katoms.end(), katoms.begin(), katoms.end());
    extract_machine_definitions();
    extract_klammer_definitions();
    m_sources.push_back(pathname);
}

void Machine::read(const std::string& s)
{
    (void)K::log(3, s);
    m_state.open_frame("Machine state: " + s);
    std::string text = m_state.subst(trim_right(s));
    katom_list katoms = process(text, command_pathname);
    m_katoms.insert(m_katoms.end(), katoms.begin(), katoms.end());
    extract_machine_definitions();
    extract_klammer_definitions();
    m_sources.push_back(s);
}    

// Read

void Machine::expand_read_katoms(
    katom_list& katoms, std::string current_filename,
    bool nonascii, bool literal, bool ignore, bool whitespace,
    bool klammers, bool eval, bool cond, bool read)
{
    (void)K::log(3);
    current_filename = resolve_relative_to(current_filename);
    // msg() << "current_filename: " << current_filename << "\n";
    if (std::find_if(katoms.begin(), katoms.end(), begin_read) != katoms.end()) {
        (void)K::log(3);
        for (const auto& [op, cl] : find_spans(katoms, begin_apply, end_apply, true, "read")) {
            auto [begin, end] = find_span_katoms(katoms, op, cl);
            if (begin_read(*begin)) {
                std::string read_filename = to_string(begin + 1, end - 1, true);

                // msg() << "read: " << resolve_relative_to(read_filename, current_filename) << "\n";

                /*
                std::string current_directory =
                    fs::path(current_filename).parent_path().string();
                fs::path input_filename =
                    fs::path(current_directory + "/" + read_filename);
                */

                fs::path input_filename = resolve_relative_to(read_filename, current_filename);
                // msg() << "read: " << input_filename << "\n";

                (void)K::log(2, input_filename.string());
                if (!fs::exists(input_filename)) {
                    std::stringstream ss{};
                    ss <<"File " << input_filename << " does not exist";
                    throw File_error(ss.str(), begin->m_loc);
                }
                std::for_each(begin, end, mark_as_replaced);

                input_filename = fs::canonical(input_filename);
                std::string text = trim_right(string_from_file(input_filename.string()));
                katom_list ks = katomize(line_split(text), input_filename);
                // ks = 
                process_katoms(
                    // ks, command_pathname,
                    ks, input_filename,
                    nonascii, literal, ignore, whitespace, klammers, eval, cond, read);
                katoms.insert(end, ks.begin(), ks.end());
            }
        }
    }
}


void Machine::extract_machine_definitions()
{
    (void)K::log(3);
    if (m_katoms.empty()) {
        return;
    }
    for (const auto& [op, cl] : find_spans(m_katoms, begin_machine_def, end_machine_def, true, command_name)) {
        auto [begin, end] = find_span_katoms(m_katoms, op, cl);
        //std::string name = trim_char(begin->m_text, '@');
        std::string name = begin->m_text;
        if (name == "@@@target") {
            m_targets.add(begin, end, m_katoms);
        } else if (name == "@@@argtype") {
            m_argtypes.add(begin, end, m_katoms);
        } else if (name == "@@@state") {
            m_state.parse_state_katoms(begin, end, m_katoms);
        }
    }
}

void Machine::extract_klammer_definitions(katom_list katoms)
{
    fmsg() << katoms << "\n";
    (void)K::log(3);
    for (const auto& [op, cl] : find_spans(katoms, begin_klammer_def, end_klammer_def, true, command_name)) {
        expand_constant_klammers(katoms, op, cl);
        auto [begin, end] = find_span_katoms(katoms, op, cl);
        m_klammers.add(m_argtypes, m_targets, begin, end, katoms);
    }
    m_klammers.rationalize(m_targets);
}

void Machine::extract_klammer_definitions()
{
    (void)K::log(3);
    for (const auto& [op, cl] : find_spans(m_katoms, begin_klammer_def, end_klammer_def, true, command_name)) {
        expand_constant_klammers(m_katoms, op, cl);
        auto [begin, end] = find_span_katoms(m_katoms, op, cl);
        m_klammers.add(m_argtypes, m_targets, begin, end, m_katoms);
    }
    m_klammers.rationalize(m_targets);
}

void Machine::update_state(const std::map<std::string, std::string>& arg_map)
{
    for (const auto& [k, v] : arg_map) {
        m_state.set(k, v);
    }
}

katom_list Machine::apply_klammer(
    Klammer& klammer, const std::string& target, katom_iter arguments_begin, katom_iter arguments_end)
{
    (void)K::log(3, "argument substitution", *arguments_begin, *(arguments_end - 1));
    m_state.replace("K_loc", arguments_begin->m_loc.str(), false);
    auto [positional, optional, rest] =
        argument_split(arguments_begin + 1, arguments_end - 1, klammer.m_parameters.m_positional.size());
    auto values = klammer.m_parameters.value_map(positional, optional, rest, arguments_begin->m_loc);

    // Resolve KTESC markers in argument values so that @eval code receives
    // the original characters (e.g., filenames with underscores).  The markers
    // remain in the klammer body substitution for final target-specific output.
    katom_list result(klammer.m_body[target].begin(), klammer.m_body[target].end());
    auto varmap = klammer.m_varmap[target];
    m_state.open_frame("Arguments for klammer " + q_(klammer.m_name));
    m_state.set(values, klammer.m_parameters);
    for (const auto& [name, indices] : varmap) {
        std::regex arg("\\*" + name + "\\*");
        for (auto i : indices) {
            result[i].m_text = std::regex_replace(result[i].m_text, arg, m_state.value(name));
            result[i].m_type = katom_t::text;
        }
    }
    // Escape target-specific characters (e.g. tex "&" -> "\&") in the writer
    // text of a GENERAL klammer's body.  Runs BEFORE process_katoms/apply()
    // below expand the body, so that target-native markup pulled in by nested
    // klammers (e.g. nl.tex -> "\newline") is left untouched -- only this
    // klammer's own literal writer text is escaped here; nested klammers escape
    // theirs when they are applied in turn.  Bodies from target-specific
    // definitions (m_body_generic[target] == false) are already in target form
    // and skipped.  KTESC markers are idempotent, so text already escaped at the
    // top level passes through unchanged.  Two kinds of body content are NOT
    // writer text and must be skipped:
    //   * ^'...'^ literal spans -- raw target markup the writer typed directly.
    //     At this point they are typed literal_begin/literal_end with plain-text
    //     content (the literal phase runs in process_katoms, below), so track
    //     span depth rather than testing katom type.
    //   * @eval / @read / @cond argument spans -- code, filenames, and
    //     predicates consumed by the primitive, NOT emitted as target text.
    //     (Escaping an underscore in "offer.Price_list(K)" broke @eval.)  The
    //     primitive's KLAMMERTEXT result, produced by process_katoms below, is
    //     klammer output and is likewise never escaped -- it is inserted after
    //     this pass and so is untouched, matching the top-level behavior where
    //     @eval is resolved before the escape pass runs.
    auto gen = klammer.m_body_generic.find(target);
    if (gen != klammer.m_body_generic.end() && gen->second) {
        Target tgt = m_targets.get(target, Locator());
        if (!tgt.m_escapes.empty()) {
            int literal_depth = 0;
            int code_depth = 0;              // inside an @eval/@read/@cond span
            std::vector<bool> apply_is_code; // one entry per open application
            for (auto& k : result) {
                if (k.m_type == katom_t::literal_begin) { ++literal_depth; continue; }
                if (k.m_type == katom_t::literal_end) {
                    if (literal_depth > 0) --literal_depth;
                    continue;
                }
                if (k.m_type == katom_t::eval_begin ||
                    k.m_type == katom_t::read_begin ||
                    k.m_type == katom_t::cond_begin) {
                    apply_is_code.push_back(true);
                    ++code_depth;
                    continue;
                }
                if (k.m_type == katom_t::apply_begin) {
                    apply_is_code.push_back(false);
                    continue;
                }
                if (k.m_type == katom_t::apply_end) {
                    if (!apply_is_code.empty()) {
                        if (apply_is_code.back()) --code_depth;
                        apply_is_code.pop_back();
                    }
                    continue;
                }
                if (literal_depth == 0 && code_depth == 0 &&
                    (k.m_type == katom_t::text ||
                     k.m_type == katom_t::word ||
                     k.m_type == katom_t::newline))
                    k.m_text = tgt.escape_text(k.m_text);
            }
        }
    }
    process_katoms(result, klammer.m_name);
    apply(m_klammers, result, target);
    m_state.close_frame();
    // msg() << boldblack << "APPLY: " << std::pair(arguments_begin, arguments_end) << "\n"
    //      << boldblack << "RESULT: " << ktype << result << black << "\n";

    modify_type(katom_t::replaced, arguments_begin, arguments_end);
    return result;
}

void Machine::apply_klammer_set(
    Klammer_set& klammer_set, katom_list& katoms, const std::string& target, katom_iter begin, katom_iter end)
{
    (void)K::log(3, "Klammer");
    std::string name = trim_char(begin->m_text, '@');
    katom_list applied_katoms = apply_klammer(klammer_set.m_klammers[name], target, begin, end);
    for (auto& k : applied_katoms) {
        if (k.m_type == katom_t::bar || k.m_type == katom_t::double_bar || k.m_type == katom_t::option_name) {
            k.m_type = katom_t::text;
        }
    }
    katoms.insert(end, applied_katoms.begin(), applied_katoms.end());
}

void Machine::apply(
    Klammer_set& klammer_set, katom_list& katoms, const std::string& target)
{
    (void)K::log(3, "Klammer_set");
    for (const auto& [op, cl] : find_spans(
             katoms, begin_klammer_apply, end_klammer_apply, true, command_name)) {
        auto [begin, end] = find_span_katoms(katoms, op, cl);
        klammer_set.check_klammer(
            klammer_name_from_katom(begin->m_text, begin->m_loc),
            target, begin->m_loc);
        apply_klammer_set(klammer_set, katoms, target, begin, end);
    }
}

std::string Machine::run_phase_functions()
{
    Target target = m_targets.get(m_state.value("K_target"), Locator());
    if (!target.m_after_apply.empty()) {
        (void)K::log(2, target);
        for (auto f : target.m_after_apply) {
            // A mode-tagged spec (":cpp ...") names a function that receives
            // the Machine itself; a bare Python function is called with the
            // result text. The Eval is constructed per phase so a chained
            // phase sees its predecessor's result in K_result.
            Eval E(*this, Locator());
            if (!f.empty() && f[0] != ':') {
                f += "(K_result)";
            }
            f = "@eval " + f + " @";
            auto katoms = katomize(line_split(f), "phase");
            // A phase function's input and output are final target text, not
            // Klammertext: take the raw result string.  Re-reading it as
            // Klammertext (Eval::eval) would misparse target characters --
            // e.g. a "@" from a quoted ^@ in justified txt output.
            m_result = E.eval_command(katoms.begin(), katoms.end() - 2);
        }
    }
    return m_result;
}


void Machine::escape_target_characters(const Target& target, katom_list& katoms)
{
    if (target.m_escapes.empty()) return;
    for (auto& k : katoms) {
        // Only escape writer content katoms — text, words, and newlines.
        // Skip structural katoms (option names, bars, klammer delimiters)
        // whose text is Klammertext syntax, not writer content.
        if (k.m_type == katom_t::text ||
            k.m_type == katom_t::word ||
            k.m_type == katom_t::newline) {
            k.m_text = target.escape_text(k.m_text);
        }
    }
}

std::string Machine::apply(const std::string& target_name, bool final_processing, bool escape_characters)
{
    (void)K::log(3, "top level");
    int recursive_limit = 5;
    m_state.set("K_target", target_name);
    m_state.subst(m_katoms.begin(), m_katoms.end());

    // Escape target-specific characters in writer text before klammer application.
    // Characters produced later by klammer bodies will not be escaped.
    // Skipped for sub-Machine apply() calls (e.g., from @eval), where the
    // text is already in target-specific form.
    auto target = m_targets.get(target_name, Locator());
    if (escape_characters)
        escape_target_characters(target, m_katoms);

    int apply_count = 0;
    auto katom_size = m_katoms.size();
    while (true) {
        apply(m_klammers, m_katoms, target_name);

        if (m_katoms.size() == katom_size) {
            break;
        }
        if (++apply_count > recursive_limit) {
            msg() << red << "Error: Recursive limit ("
                  << recursive_limit << ") reached\n" << black;
            break;
        }
        katom_size = m_katoms.size();
    }
    m_result = to_string(m_katoms.begin(), m_katoms.end());

    if (final_processing) {
        for (const auto& [old_str, new_str] : target.m_transforms) {
            m_result = string_replace(m_result, old_str, new_str);
        }
        m_result = target.resolve_escapes(m_result);
        m_result = run_phase_functions();
    }
    m_result = trim_char(m_result, '\n');
    return m_result;
}    
