#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdio>
#include <cstring>
#include <expected>
#include <fstream>
#include <memory>
#include <optional>
#include <print>
#include <sstream>
#include <stdexcept>
#include <string_view>
#include <unordered_map>
#include <variant>
#include <vector>

using std::size_t;
using namespace std::literals;

struct Loc {
    size_t idx;
    int line;
    int col;
};

struct File {
    std::string filename;
    std::string text;

    static auto read_text_file(std::string_view filename)
        -> std::expected<File, std::string>
    {
        auto file_stream = std::ifstream(filename.data());
        if (!file_stream) {
            return std::unexpected("could not open file");
        }

        auto text_stream = std::stringstream();
        text_stream << file_stream.rdbuf();
        auto text = text_stream.str();

        return File {
            .filename = std::string(filename),
            .text = std::move(text),
        };
    }
};

void print_error(File& file, Loc loc, std::string_view message)
{
    auto& text = file.text;
    auto [idx, line, col] = loc;

    constexpr auto type = "error"sv;

    constexpr auto clear = "\x1b[0m"sv;
    constexpr auto bold_red = "\x1b[1;91m"sv;
    constexpr auto bold_white = "\x1b[1;97m"sv;
    constexpr auto cyan = "\x1b[0;36m"sv;
    constexpr auto gray = "\x1b[0;90m"sv;
    constexpr auto light_gray = "\x1b[0;37m"sv;
    constexpr auto green = "\x1b[0;32m"sv;

    auto start = text.find_last_of('\n', idx) + 1;
    auto end = text.find_first_of('\n', idx);

    if (end == std::string_view::npos) {
        end = text.size();
    }

    auto line_text = text.substr(start, end - start);

    auto linenr_str = std::to_string(line);

    // clang-format off
    std::println(""
        "{0}{1}{2}: {3}\n"
        " {4}--> {5}:{6}:{7}\n"
        " {8: <{9}}{10}|\n"
        " {11}{12}{13}|{14}{15}\n"
        " {16: <{17}}"
        "{18}|{19: <{20}}{21}^ {22}{23}{24}\n",
        bold_red,          type,       bold_white,
        message,           cyan,       file.filename,
        line,              col,        "",
        linenr_str.size(), gray,       light_gray,
        linenr_str,        gray,       green,
        line_text,         "",         linenr_str.size(),
        gray,              "",         col - 1,
        bold_red,          bold_white, message,
        clear);
    // clang-format on
}

enum class TokTy {
    Ident,
    Str,
    Int = '0',
    LParen = '(',
    RParen = ')',
};

struct Tok {
    TokTy ty;
    Loc loc;
    std::string_view text;
};

auto tokenize(File& file) -> std::vector<Tok>
{
    auto& text = file.text;

    size_t i = 0;
    int line = 1;
    int col = 1;

    auto toks = std::vector<Tok>();

    bool error_occured = false;

    auto step = [&]() {
        if (text[i] == '\n') {
            line += 1;
            col = 1;
        } else {
            col += 1;
        }
        i += 1;
    };

    auto push = [&](Loc loc, TokTy ty) {
        toks.push_back(Tok {
            .ty = ty,
            .loc = loc,
            .text = std::string_view(&text[loc.idx], i - loc.idx),
        });
    };

    while (i < text.size()) {
        auto loc = Loc { i, line, col };

        bool matched = false;
        while (i < text.size() && std::strchr(" \t\r\n", text[i])) {
            matched = true;
            step();
        }
        if (matched)
            continue;
        while (i < text.size()
            && ((text[i] >= 'a' && text[i] <= 'z')
                || (text[i] >= 'A' && text[i] <= 'Z')
                || (matched && text[i] >= '0' && text[i] <= '9')
                || std::strchr("_+-*/%=<>", text[i]))) {
            matched = true;
            step();
        }
        if (matched) {
            push(loc, TokTy::Ident);
            continue;
        }
        while (i < text.size()
            && ((text[i] >= '1' && text[i] <= '9')
                || (matched && text[i] == '0'))) {
            matched = true;
            step();
        }
        if (matched) {
            push(loc, TokTy::Int);
            continue;
        }
        if (std::strchr("()0", text[i])) {
            auto ty = TokTy(text[i]);
            step();
            push(loc, ty);
            continue;
        }
        if (text[i] == '"') {
            step();
            while (i < text.size() && text[i] != '"') {
                if (text[i] == '\\') {
                    step();
                    if (i >= text.size())
                        break;
                }
                step();
            }
            if (i >= text.size() || text[i] != '"') {
                print_error(file, loc, "malformed string");
                error_occured = true;
                continue;
            }
            step();
            push(loc, TokTy::Str);
            continue;
        }
        print_error(file, loc, "illegal character");
        step();
        error_occured = true;
    }

    if (error_occured) {
        std::println(stderr, "error(s) occured. aborting.");
        std::exit(1);
    }

    return toks;
}

enum class SExprTy {
    Ident,
    Int,
    Str,
    List,
};

class SExpr {
private:
    struct PrivKey { };

public:
    explicit SExpr([[maybe_unused]] PrivKey key,
        SExprTy ty,
        Loc loc,
        std::variant<Tok, std::vector<std::unique_ptr<SExpr>>> data)
        : m_ty(ty)
        , m_loc(loc)
        , m_data(std::move(data)) { };

    static auto make_ident(Loc loc, Tok tok) -> std::unique_ptr<SExpr>
    {
        return std::make_unique<SExpr>(PrivKey { }, SExprTy::Ident, loc, tok);
    }
    static auto make_int(Loc loc, Tok tok) -> std::unique_ptr<SExpr>
    {
        return std::make_unique<SExpr>(PrivKey { }, SExprTy::Int, loc, tok);
    }
    static auto make_str(Loc loc, Tok tok) -> std::unique_ptr<SExpr>
    {
        return std::make_unique<SExpr>(PrivKey { }, SExprTy::Str, loc, tok);
    }
    static auto make_list(Loc loc, std::vector<std::unique_ptr<SExpr>> exprs)
        -> std::unique_ptr<SExpr>
    {
        return std::make_unique<SExpr>(
            PrivKey { }, SExprTy::List, loc, std::move(exprs));
    }

    auto ty() const -> SExprTy
    {
        return m_ty;
    }

    auto tok() const -> Tok
    {
        assert(m_ty == SExprTy::Ident || m_ty == SExprTy::Int
            || m_ty == SExprTy::Str);
        return std::get<Tok>(m_data);
    }

    auto nodes() const -> const std::vector<std::unique_ptr<SExpr>>&
    {
        assert(m_ty == SExprTy::List);
        return std::get<std::vector<std::unique_ptr<SExpr>>>(m_data);
    }

    void print() const
    {
        if (m_ty == SExprTy::List) {
            std::print("(");
            bool first = true;
            for (auto& node : nodes()) {
                if (not first) {
                    std::print(" ");
                }
                first = false;
                node->print();
            }
            std::print(")");
        } else {
            std::print("{}", tok().text);
        }
    }

private:
    SExprTy m_ty;
    Loc m_loc;
    std::variant<Tok, std::vector<std::unique_ptr<SExpr>>> m_data;
};

auto parse_sexprs(File& file) -> std::unique_ptr<SExpr>
{
    auto toks = tokenize(file);
    auto tok = toks.begin();

    auto begin_loc = tok->loc;

    auto parse_expr = [&]() -> std::unique_ptr<SExpr> {
        auto impl = [&](auto& self) -> std::unique_ptr<SExpr> {
            auto loc = tok->loc;
            if (tok->ty == TokTy::Ident) {
                auto t = *tok;
                ++tok;
                return SExpr::make_ident(loc, t);
            } else if (tok->ty == TokTy::Int) {
                auto t = *tok;
                ++tok;
                return SExpr::make_int(loc, t);
            } else if (tok->ty == TokTy::Str) {
                auto t = *tok;
                ++tok;
                return SExpr::make_str(loc, t);
            } else if (tok->ty == TokTy::LParen) {
                ++tok;
                auto exprs = std::vector<std::unique_ptr<SExpr>>();
                while (tok != toks.end() && tok->ty != TokTy::RParen) {
                    auto expr = self(self);
                    if (!expr)
                        return nullptr;
                    exprs.push_back(std::move(expr));
                }
                if (tok == toks.end() || tok->ty != TokTy::RParen) {
                    print_error(file,
                        tok != toks.end() ? tok->loc : loc,
                        "expected ')'");
                    return nullptr;
                }
                ++tok;
                return SExpr::make_list(loc, std::move(exprs));
            } else {
                print_error(file, loc, "expected expression");
                return nullptr;
            }
        };
        return impl(impl);
    };

    auto exprs = std::vector<std::unique_ptr<SExpr>>();
    while (tok != toks.end()) {
        auto expr = parse_expr();
        if (!expr)
            return nullptr;
        exprs.push_back(std::move(expr));
    }
    return SExpr::make_list(begin_loc, std::move(exprs));
}

class SEMatcher {
public:
    using SEPtr = std::unique_ptr<SExpr>;

    explicit SEMatcher(SExpr& node)
        : m_node(node) { };

    auto match() -> std::optional<std::unordered_map<std::string, SExpr*>>
    {
        return not m_failed ? std::optional(std::move(m_captures))
                            : std::nullopt;
    }

    auto list() -> SEMatcher&
    {
        if (m_node.ty() != SExprTy::List) {
            m_failed = true;
        }
        return *this;
    }

    auto list(size_t count) -> SEMatcher&
    {
        if (m_node.ty() != SExprTy::List || m_node.nodes().size() != count) {
            m_failed = true;
        }
        return *this;
    }

    auto at(size_t idx, auto func) -> SEMatcher&
    {
        if (m_node.ty() != SExprTy::List || idx >= m_node.nodes().size()) {
            m_failed = true;
            return *this;
        }
        auto inner = SEMatcher(*m_node.nodes().at(idx));
        func(inner);
        m_captures.merge(inner.m_captures);
        m_failed |= inner.m_failed;
        return *this;
    }

    auto all(auto func) -> SEMatcher&
    {
        if (m_node.ty() != SExprTy::List) {
            m_failed = true;
            return *this;
        }
        for (auto& inner_node : m_node.nodes()) {
            auto inner = SEMatcher(*inner_node);
            func(inner);
            m_captures.merge(inner.m_captures);
            m_failed |= inner.m_failed;
        }
        return *this;
    }

    auto capture(std::string id) -> SEMatcher&
    {
        m_captures[id] = &m_node;
        return *this;
    }

    auto ident() -> SEMatcher&
    {
        if (m_node.ty() != SExprTy::Ident) {
            m_failed = true;
        }
        return *this;
    }

    auto ident(std::string val) -> SEMatcher&
    {
        if (m_node.ty() != SExprTy::Ident || m_node.tok().text != val) {
            m_failed = true;
        }
        return *this;
    }

private:
    SExpr& m_node;
    bool m_failed = false;
    std::unordered_map<std::string, SExpr*> m_captures = { };
};

enum NodeTy {
    File,
    FnStmt,
    CallExpr,
};

class Node {
private:
    struct PrivKey { };

public:
    using Vec = std::vector<std::unique_ptr<Node>>;

    struct File {
        Vec stmts;
    };

    struct FnStmt {
        std::string ident;
        Vec params;
        std::unique_ptr<Node> ret_ty;
        Vec body;
    };

    using Data = std::variant<File>;

    explicit Node([[maybe_unused]] PrivKey key, NodeTy ty, Data data)
        : m_ty(ty)
        , m_data(std::move(data)) { };

    static auto make_file(File file) -> std::unique_ptr<Node>
    {
        return std::make_unique<Node>(
            PrivKey { }, NodeTy::File, std::move(file));
    }

private:
    NodeTy m_ty;
    Data m_data;
};

class Parser {
public:
    using M = SEMatcher;
    using SEPtr = std::unique_ptr<SExpr>;
    using NodePtr = std::unique_ptr<Node>;

    auto parse_file(SExpr& node) -> NodePtr
    {
        auto stmts = Node::Vec();
        for (auto& inner : node.nodes()) {
            stmts.push_back(parse_stmt(*inner));
        }
        return Node::make_file(Node::File { std::move(stmts) });
    }

    auto parse_stmt(SExpr& node) -> NodePtr
    {
        auto fn_stmt_match
            = M(node)
                  .list(5)
                  .at(0, [](M& m) { return m.ident("fn"); })
                  .at(1, [](M& m) { return m.ident().capture("ident"); })
                  .at(2, [](M& m) { return m.list().capture("params"); })
                  .at(3, [](M& m) { return m.capture("retty"); })
                  .at(4, [](M& m) { return m.list().capture("body"); })
                  .match();
        if (fn_stmt_match) {
            auto m = std::move(*fn_stmt_match);
            std::println("matched!, ident = {}", m["ident"]->tok().text);
        }
        std::println("not matched :(");
        return nullptr;
    }
};

int main(int argc, char* argv[])
{
    assert(argc >= 2);
    auto file = File::read_text_file(argv[1]).value();
    auto sexpr_ast = parse_sexprs(file);
    std::println();

    auto file_ast = Parser().parse_file(*sexpr_ast);
}