#include "cf.h"

#include "str.h"

#include <stdio.h>
#include <string.h>

static int has_cfg_note(const tobi_ir *n, size_t off) {
    for (size_t i = 0; i < n->child_len; i++) {
        const tobi_ir *c = n->child[i];
        if (c->kind == TOBI_IR_DIAG && c->off == off && c->str && strstr(c->str, "cfg fallback") != NULL) {
            return 1;
        }
    }
    return 0;
}

static int unclear_control(const tobi_ir *n) {
    if (n->kind != TOBI_IR_UNKNOWN) {
        return 0;
    }
    return n->raw_op == 0xa1u || n->raw_op == 0xccu || n->raw_op == 0x5b31u || n->raw_op == 0x5b32u || n->raw_op == 0u;
}

static void add_cfg_note(tobi_ir *parent, const tobi_ir *unknown) {
    if (!parent || !unknown || has_cfg_note(parent, unknown->off)) {
        return;
    }
    char msg[128];
    (void)snprintf(msg, sizeof(msg), "cfg fallback label L_%04zx for unclear opcode 0x%x", unknown->off, unknown->raw_op);
    tobi_ir *note = tobi_ir_new(TOBI_IR_DIAG, unknown->off, unknown->len);
    note->raw_op = unknown->raw_op;
    tobi_ir_set_name(note, "cfg");
    tobi_ir_set_str(note, msg);
    tobi_ir_add(parent, note);
}

static void add_diag_node(tobi_diag_list *diag, tobi_diag_level level, const tobi_ir *n, const char *msg) {
    tobi_diag_add_source(diag, level, n ? n->off : 0, n ? n->source : NULL,
                         n ? n->input_index : 0, n ? n->has_source : 0, "%s", msg);
}

static void walk(tobi_ir *n, tobi_diag_list *diag, unsigned loop_depth) {
    if (!n) {
        return;
    }
    if (n->kind == TOBI_IR_IF && n->child_len < 2) {
        add_diag_node(diag, TOBI_DIAG_ERROR, n, "if node lacks predicate or body");
    }
    if (n->kind == TOBI_IR_WHILE && n->child_len < 2) {
        add_diag_node(diag, TOBI_DIAG_ERROR, n, "while node lacks predicate or body");
    }
    if ((n->kind == TOBI_IR_BREAK || n->kind == TOBI_IR_CONTINUE) && loop_depth == 0) {
        tobi_diag_add_source(diag, TOBI_DIAG_WARN, n->off, n->source, n->input_index, n->has_source,
                             "%s outside loop kept as explicit statement",
                             n->kind == TOBI_IR_BREAK ? "break" : "continue");
    }
    size_t original_children = n->child_len;
    unsigned child_loop_depth = loop_depth + (n->kind == TOBI_IR_WHILE ? 1u : 0u);
    for (size_t i = 0; i < original_children; i++) {
        if ((n->kind == TOBI_IR_BLOCK || n->kind == TOBI_IR_ROOT) && unclear_control(n->child[i])) {
            add_cfg_note(n, n->child[i]);
        }
        walk(n->child[i], diag, child_loop_depth);
    }
}

int tobi_cf_recover(tobi_ir *root, tobi_diag_list *diag) {
    walk(root, diag, 0);
    tobi_ir_assign_ids(root);
    return !tobi_diag_has_error(diag);
}

static void dot_escape_limited(tobi_sb *sb, const char *s, size_t wrap, size_t max) {
    if (!s) {
        return;
    }
    size_t col = 0;
    size_t emitted = 0;
    for (const char *p = s; *p; p++) {
        if (max && emitted >= max) {
            tobi_sb_add(sb, "\\n...");
            return;
        }
        unsigned char c = (unsigned char)*p;
        if (c == '\n' || c == '\r') {
            tobi_sb_add(sb, "\\n");
            col = 0;
            continue;
        }
        if (wrap && col >= wrap) {
            tobi_sb_add(sb, "\\n");
            col = 0;
        }
        if (c == '"' || c == '\\') {
            tobi_sb_ch(sb, '\\');
        }
        if (c < 0x20) {
            tobi_sb_ch(sb, ' ');
        } else {
            tobi_sb_ch(sb, (char)c);
        }
        col++;
        emitted++;
    }
}

static void dot_escape(tobi_sb *sb, const char *s) {
    dot_escape_limited(sb, s, 0, 0);
}

static void dot_label_text(tobi_sb *sb, const char *s) {
    dot_escape_limited(sb, s, 28, 96);
}

typedef struct {
    size_t stmt_next;
    size_t edge_next;
} dot_ctx;

static size_t dot_stmt(tobi_sb *sb, const tobi_ir *n, dot_ctx *ctx);

static void dot_edge(tobi_sb *sb, dot_ctx *ctx, size_t from, size_t to, const char *label) {
    if (label && label[0]) {
        size_t id = ctx->edge_next++;
        tobi_sb_printf(sb, "  e%zu [shape=box,style=\"solid\",", id);
        tobi_sb_add(sb, "color=\"#000000\",fontcolor=\"#000000\",");
        tobi_sb_add(sb, "fontname=\"Times-Roman\",fontsize=10,margin=\"0.06,0.03\",");
        tobi_sb_add(sb, "height=0,width=0,label=\"");
        dot_escape(sb, label);
        tobi_sb_add(sb, "\"];\n");
        tobi_sb_printf(sb, "  n%zu -> e%zu [arrowhead=none, weight=2];\n", from, id);
        tobi_sb_printf(sb, "  e%zu -> n%zu [weight=2];\n", id, to);
        return;
    }
    tobi_sb_printf(sb, "  n%zu -> n%zu;\n", from, to);
}

static int dot_block(tobi_sb *sb, const tobi_ir *b, size_t owner, const char *edge_label, dot_ctx *ctx, size_t *last) {
    if (!b) {
        return 0;
    }
    size_t prev = owner;
    int emitted = 0;
    for (size_t i = 0; i < b->child_len; i++) {
        size_t id = dot_stmt(sb, b->child[i], ctx);
        dot_edge(sb, ctx, prev, id, i == 0 ? edge_label : "next");
        prev = id;
        emitted = 1;
    }
    if (last) {
        *last = prev;
    }
    return emitted;
}

static void dot_ir_label(tobi_sb *sb, const tobi_ir *n) {
    tobi_sb_add(sb, tobi_ir_kind_name(n->kind));
    if (n->name && n->name[0]) {
        tobi_sb_add(sb, "\\n");
        dot_label_text(sb, n->name);
    } else if (n->path && n->path[0]) {
        tobi_sb_add(sb, "\\n");
        dot_label_text(sb, n->path);
    } else if (n->str && n->str[0]) {
        tobi_sb_add(sb, "\\n");
        dot_label_text(sb, n->str);
    }
    tobi_sb_printf(sb, "\\noff=0x%zx", n->off);
    if (n->kind == TOBI_IR_UNKNOWN) {
        tobi_sb_printf(sb, "\\nop=0x%x", n->raw_op);
    }
}

static size_t dot_stmt(tobi_sb *sb, const tobi_ir *n, dot_ctx *ctx) {
    size_t id = ctx->stmt_next++;
    tobi_sb_printf(sb, "  n%zu [label=\"", id);
    dot_ir_label(sb, n);
    tobi_sb_add(sb, "\"];\n");
    if (n->kind == TOBI_IR_METHOD && n->child_len > 0) {
        (void)dot_block(sb, n->child[0], id, "entry", ctx, NULL);
    } else if (n->kind == TOBI_IR_SCOPE || n->kind == TOBI_IR_DEVICE || n->kind == TOBI_IR_PROCESSOR) {
        if (n->child_len > 0) {
            (void)dot_block(sb, n->child[0], id, "body", ctx, NULL);
        }
    } else if (n->kind == TOBI_IR_IF) {
        if (n->child_len > 1) {
            (void)dot_block(sb, n->child[1], id, "then", ctx, NULL);
        }
        if (n->child_len > 2) {
            (void)dot_block(sb, n->child[2], id, "else", ctx, NULL);
        }
    } else if (n->kind == TOBI_IR_WHILE) {
        if (n->child_len > 1) {
            size_t last = id;
            if (dot_block(sb, n->child[1], id, "while-body", ctx, &last)) {
                dot_edge(sb, ctx, last, id, "back");
            }
        }
    }
    return id;
}

static void dot_common_attrs(tobi_sb *sb) {
    tobi_sb_add(sb, "  graph [rankdir=TB,bgcolor=\"#ffffff\",pad=0.18,nodesep=0.5,ranksep=0.72,");
    tobi_sb_add(sb, "splines=ortho,outputorder=edgesfirst];\n");
    tobi_sb_add(sb, "  node [shape=box,style=\"solid\",color=\"#000000\",fontcolor=\"#000000\",");
    tobi_sb_add(sb, "fontname=\"Courier\",fontsize=12,margin=\"0.12,0.07\",penwidth=1.1];\n");
    tobi_sb_add(sb, "  edge [color=\"#000000\",fontcolor=\"#000000\",arrowsize=0.75,penwidth=1.1];\n");
}

static char *dot_cfg(const tobi_ir *root) {
    tobi_sb sb;
    tobi_sb_init(&sb);
    tobi_sb_add(&sb, "digraph tobi_cfg {\n");
    dot_common_attrs(&sb);
    if (root) {
        dot_ctx ctx = {0, 0};
        size_t root_id = dot_stmt(&sb, root, &ctx);
        size_t prev = root_id;
        for (size_t i = 0; i < root->child_len; i++) {
            size_t id = dot_stmt(&sb, root->child[i], &ctx);
            dot_edge(&sb, &ctx, prev, id, i == 0 ? "entry" : "next");
            prev = id;
        }
    }
    tobi_sb_add(&sb, "}\n");
    return tobi_sb_take(&sb);
}

static size_t dot_ast_node(tobi_sb *sb, const tobi_ir *n, dot_ctx *ctx) {
    size_t id = ctx->stmt_next++;
    tobi_sb_printf(sb, "  n%zu [label=\"", id);
    dot_ir_label(sb, n);
    tobi_sb_add(sb, "\"];\n");
    for (size_t i = 0; i < n->child_len; i++) {
        size_t child_id = dot_ast_node(sb, n->child[i], ctx);
        tobi_sb_printf(sb, "  n%zu -> n%zu;\n", id, child_id);
    }
    return id;
}

static char *dot_ast(const tobi_ir *root) {
    tobi_sb sb;
    tobi_sb_init(&sb);
    tobi_sb_add(&sb, "digraph tobi_ast {\n");
    dot_common_attrs(&sb);
    if (root) {
        dot_ctx ctx = {0, 0};
        (void)dot_ast_node(&sb, root, &ctx);
    }
    tobi_sb_add(&sb, "}\n");
    return tobi_sb_take(&sb);
}

static int ns_index(const tobi_ns *ns, const char *path, size_t *idx) {
    if (!ns || !path) {
        return 0;
    }
    for (size_t i = 0; i < ns->len; i++) {
        if (ns->items[i].path && strcmp(ns->items[i].path, path) == 0) {
            *idx = i;
            return 1;
        }
    }
    return 0;
}

static char *dot_namespace(const tobi_ns *ns) {
    tobi_sb sb;
    tobi_sb_init(&sb);
    tobi_sb_add(&sb, "digraph tobi_namespace {\n");
    dot_common_attrs(&sb);
    tobi_sb_add(&sb, "  n0 [label=\"namespace\\n\\\\\"];\n");
    if (ns) {
        for (size_t i = 0; i < ns->len; i++) {
            const tobi_ns_ent *e = &ns->items[i];
            tobi_sb_printf(&sb, "  n%zu [label=\"%s", i + 1u, tobi_ns_kind_name(e->kind));
            if (e->name && e->name[0]) {
                tobi_sb_add(&sb, "\\n");
                dot_label_text(&sb, e->name);
            }
            if (e->path && e->path[0]) {
                tobi_sb_add(&sb, "\\n");
                dot_label_text(&sb, e->path);
            }
            tobi_sb_printf(&sb, "\\noff=0x%zx\"];\n", e->off);
        }
        for (size_t i = 0; i < ns->len; i++) {
            const tobi_ns_ent *e = &ns->items[i];
            size_t parent = 0;
            size_t found = 0;
            if (e->owner && ns_index(ns, e->owner, &found)) {
                parent = found + 1u;
            }
            tobi_sb_printf(&sb, "  n%zu -> n%zu;\n", parent, i + 1u);
            if (e->target && ns_index(ns, e->target, &found)) {
                tobi_sb_printf(&sb, "  n%zu -> n%zu [style=dashed,arrowhead=vee];\n", i + 1u, found + 1u);
            }
        }
    }
    tobi_sb_add(&sb, "}\n");
    return tobi_sb_take(&sb);
}

char *tobi_cf_dot(const tobi_ir *root) {
    return dot_cfg(root);
}

char *tobi_graph_dot(const tobi_ir *root, const tobi_ns *ns, tobi_graph_kind kind) {
    switch (kind) {
    case TOBI_GRAPH_AST:
        return dot_ast(root);
    case TOBI_GRAPH_NAMESPACE:
        return dot_namespace(ns);
    case TOBI_GRAPH_CFG:
    default:
        return dot_cfg(root);
    }
}