#include "expr.h"

#include <stdlib.h>
#include <string.h>

Expr *expr_number(double value) {
    Expr *e = malloc(sizeof *e);
    if (e == NULL) {
        return NULL;
    }
    e->kind = EXPR_NUMBER;
    e->as.number = value;
    return e;
}

Expr *expr_variable(const char *name) {
    Expr *e = malloc(sizeof *e);
    if (e == NULL) {
        return NULL;
    }
    e->as.name = malloc(strlen(name) + 1);
    if (e->as.name == NULL) {
        free(e);
        return NULL;
    }
    strcpy(e->as.name, name);
    e->kind = EXPR_VARIABLE;
    return e;
}

Expr *expr_binary(BinaryOp op, Expr *left, Expr *right) {
    Expr *e = NULL;
    if (left != NULL && right != NULL) {
        e = malloc(sizeof *e);
    }
    if (e == NULL) {
        expr_free(left);
        expr_free(right);
        return NULL;
    }
    e->kind = EXPR_BINARY;
    e->as.binary.op = op;
    e->as.binary.left = left;
    e->as.binary.right = right;
    return e;
}

Expr *expr_if(Expr *condition, Expr *then_branch, Expr *else_branch) {
    Expr *e = NULL;
    if (condition != NULL && then_branch != NULL && else_branch != NULL) {
        e = malloc(sizeof *e);
    }
    if (e == NULL) {
        expr_free(condition);
        expr_free(then_branch);
        expr_free(else_branch);
        return NULL;
    }
    e->kind = EXPR_IF;
    e->as.conditional.condition = condition;
    e->as.conditional.then_branch = then_branch;
    e->as.conditional.else_branch = else_branch;
    return e;
}

static void free_args(Expr **args, int arg_count) {
    for (int i = 0; i < arg_count; i++) {
        expr_free(args[i]);
    }
    free(args);
}

Expr *expr_call(const char *name, Expr **args, int arg_count) {
    int complete = 1;
    for (int i = 0; i < arg_count; i++) {
        if (args[i] == NULL) {
            complete = 0;
        }
    }
    Expr *e = complete ? malloc(sizeof *e) : NULL;
    char *copy = complete ? malloc(strlen(name) + 1) : NULL;
    if (e == NULL || copy == NULL) {
        free(e);
        free(copy);
        free_args(args, arg_count);
        return NULL;
    }
    strcpy(copy, name);
    e->kind = EXPR_CALL;
    e->as.call.name = copy;
    e->as.call.args = args;
    e->as.call.arg_count = arg_count;
    return e;
}

void expr_free(Expr *e) {
    if (e == NULL) {
        return;
    }
    switch (e->kind) {
        case EXPR_NUMBER:
            break;
        case EXPR_VARIABLE:
            free(e->as.name);
            break;
        case EXPR_BINARY:
            expr_free(e->as.binary.left);
            expr_free(e->as.binary.right);
            break;
        case EXPR_IF:
            expr_free(e->as.conditional.condition);
            expr_free(e->as.conditional.then_branch);
            expr_free(e->as.conditional.else_branch);
            break;
        case EXPR_CALL:
            free(e->as.call.name);
            free_args(e->as.call.args, e->as.call.arg_count);
            break;
    }
    free(e);
}

static const char *op_symbol(BinaryOp op) {
    switch (op) {
        case OP_ADD: return "+";
        case OP_SUB: return "-";
        case OP_MUL: return "*";
        case OP_DIV: return "/";
        case OP_LESS: return "<";
        case OP_LESS_EQUAL: return "<=";
        case OP_GREATER: return ">";
        case OP_GREATER_EQUAL: return ">=";
        case OP_EQUAL: return "==";
        case OP_NOT_EQUAL: return "!=";
    }
    return "?";
}

// A mayor número, más fuerte: el nivel de la gramática.
static int precedence(BinaryOp op) {
    switch (op) {
        case OP_MUL:
        case OP_DIV:
            return 3;
        case OP_ADD:
        case OP_SUB:
            return 2;
        case OP_LESS:
        case OP_LESS_EQUAL:
        case OP_GREATER:
        case OP_GREATER_EQUAL:
        case OP_EQUAL:
        case OP_NOT_EQUAL:
            return 1;
    }
    return 0;
}

// Escribe un operando, entre paréntesis si su nivel es menor que minimum. Un if siempre los
// lleva: en la gramática, solo entra en una operación con paréntesis.
static void print_child(const Expr *child, int minimum, FILE *out) {
    int needs_parens = child->kind == EXPR_IF
        || (child->kind == EXPR_BINARY && precedence(child->as.binary.op) < minimum);
    if (needs_parens) {
        fputc('(', out);
    }
    expr_print(child, out);
    if (needs_parens) {
        fputc(')', out);
    }
}

void expr_print(const Expr *e, FILE *out) {
    switch (e->kind) {
        case EXPR_NUMBER:
            fprintf(out, "%g", e->as.number);
            break;
        case EXPR_VARIABLE:
            fputs(e->as.name, out);
            break;
        case EXPR_BINARY: {
            int level = precedence(e->as.binary.op);
            // Asocian por la izquierda: a la derecha, el mismo nivel lleva paréntesis.
            // Las comparaciones no se encadenan: ahí los lleva a los dos lados.
            print_child(e->as.binary.left, level == 1 ? 2 : level, out);
            fprintf(out, " %s ", op_symbol(e->as.binary.op));
            print_child(e->as.binary.right, level + 1, out);
            break;
        }
        case EXPR_IF:
            fputs("if ", out);
            expr_print(e->as.conditional.condition, out);
            fputs(" then ", out);
            expr_print(e->as.conditional.then_branch, out);
            fputs(" else ", out);
            expr_print(e->as.conditional.else_branch, out);
            break;
        case EXPR_CALL:
            fprintf(out, "%s(", e->as.call.name);
            for (int i = 0; i < e->as.call.arg_count; i++) {
                if (i > 0) {
                    fputs(", ", out);
                }
                expr_print(e->as.call.args[i], out);
            }
            fputc(')', out);
            break;
    }
}

int expr_apply(BinaryOp op, double a, double b, double *result) {
    switch (op) {
        case OP_ADD: *result = a + b; return 1;
        case OP_SUB: *result = a - b; return 1;
        case OP_MUL: *result = a * b; return 1;
        case OP_DIV:
            if (b == 0.0) {
                return 0;
            }
            *result = a / b;
            return 1;
        case OP_LESS: *result = a < b; return 1;
        case OP_LESS_EQUAL: *result = a <= b; return 1;
        case OP_GREATER: *result = a > b; return 1;
        case OP_GREATER_EQUAL: *result = a >= b; return 1;
        case OP_EQUAL: *result = a == b; return 1;
        case OP_NOT_EQUAL: *result = a != b; return 1;
    }
    return 0;
}

int expr_eval(const Expr *e, const Env *env, double *result) {
    switch (e->kind) {
        case EXPR_NUMBER:
            *result = e->as.number;
            return 1;
        case EXPR_VARIABLE:
            return env_get(env, e->as.name, result);
        case EXPR_BINARY: {
            double left;
            double right;
            if (!expr_eval(e->as.binary.left, env, &left) || !expr_eval(e->as.binary.right, env, &right)) {
                return 0;
            }
            return expr_apply(e->as.binary.op, left, right, result);
        }
        case EXPR_IF: {
            double condition;
            if (!expr_eval(e->as.conditional.condition, env, &condition)) {
                return 0;
            }
            return expr_eval(condition != 0 ? e->as.conditional.then_branch : e->as.conditional.else_branch, env, result);
        }
        case EXPR_CALL:
            return 0;
    }
    return 0;
}

Expr *expr_copy(const Expr *e) {
    if (e == NULL) {
        return NULL;
    }
    switch (e->kind) {
        case EXPR_NUMBER:
            return expr_number(e->as.number);
        case EXPR_VARIABLE:
            return expr_variable(e->as.name);
        case EXPR_BINARY:
            return expr_binary(e->as.binary.op, expr_copy(e->as.binary.left), expr_copy(e->as.binary.right));
        case EXPR_IF:
            return expr_if(expr_copy(e->as.conditional.condition), expr_copy(e->as.conditional.then_branch),
                           expr_copy(e->as.conditional.else_branch));
        case EXPR_CALL: {
            int count = e->as.call.arg_count;
            Expr **args = NULL;
            if (count > 0) {
                args = malloc((size_t) count * sizeof *args);
                if (args == NULL) {
                    return NULL;
                }
            }
            for (int i = 0; i < count; i++) {
                args[i] = expr_copy(e->as.call.args[i]);
            }
            return expr_call(e->as.call.name, args, count);
        }
    }
    return NULL;
}

int expr_equal(const Expr *a, const Expr *b) {
    if (a->kind != b->kind) {
        return 0;
    }
    switch (a->kind) {
        case EXPR_NUMBER:
            return a->as.number == b->as.number;
        case EXPR_VARIABLE:
            return strcmp(a->as.name, b->as.name) == 0;
        case EXPR_BINARY:
            return a->as.binary.op == b->as.binary.op && expr_equal(a->as.binary.left, b->as.binary.left)
                && expr_equal(a->as.binary.right, b->as.binary.right);
        case EXPR_IF:
            return expr_equal(a->as.conditional.condition, b->as.conditional.condition)
                && expr_equal(a->as.conditional.then_branch, b->as.conditional.then_branch)
                && expr_equal(a->as.conditional.else_branch, b->as.conditional.else_branch);
        case EXPR_CALL:
            if (strcmp(a->as.call.name, b->as.call.name) != 0 || a->as.call.arg_count != b->as.call.arg_count) {
                return 0;
            }
            for (int i = 0; i < a->as.call.arg_count; i++) {
                if (!expr_equal(a->as.call.args[i], b->as.call.args[i])) {
                    return 0;
                }
            }
            return 1;
    }
    return 0;
}

static int is_number(const Expr *e, double value) {
    return e->kind == EXPR_NUMBER && e->as.number == value;
}

// Se queda con un hijo: libera el otro y el nodo. Devuelve el hijo elegido.
static Expr *keep_child(Expr *e, Expr *kept, Expr *dropped) {
    expr_free(dropped);
    free(e);
    return kept;
}

// Convierte el nodo en un número, en su sitio, y libera sus hijos.
static Expr *become_number(Expr *e, double value) {
    expr_free(e->as.binary.left);
    expr_free(e->as.binary.right);
    e->kind = EXPR_NUMBER;
    e->as.number = value;
    return e;
}

Expr *expr_simplify(Expr *e) {
    if (e == NULL || e->kind == EXPR_NUMBER || e->kind == EXPR_VARIABLE) {
        return e;
    }
    if (e->kind == EXPR_IF) {
        e->as.conditional.condition = expr_simplify(e->as.conditional.condition);
        e->as.conditional.then_branch = expr_simplify(e->as.conditional.then_branch);
        e->as.conditional.else_branch = expr_simplify(e->as.conditional.else_branch);
        return e;
    }
    if (e->kind == EXPR_CALL) {
        for (int i = 0; i < e->as.call.arg_count; i++) {
            e->as.call.args[i] = expr_simplify(e->as.call.args[i]);
        }
        return e;
    }
    e->as.binary.left = expr_simplify(e->as.binary.left);
    e->as.binary.right = expr_simplify(e->as.binary.right);
    Expr *left = e->as.binary.left;
    Expr *right = e->as.binary.right;
    BinaryOp op = e->as.binary.op;

    if ((op == OP_ADD && is_number(left, 0.0)) || (op == OP_MUL && is_number(left, 1.0))) {
        return keep_child(e, right, left);
    }
    if ((op == OP_ADD || op == OP_SUB) && is_number(right, 0.0)) {
        return keep_child(e, left, right);
    }
    if ((op == OP_MUL || op == OP_DIV) && is_number(right, 1.0)) {
        return keep_child(e, left, right);
    }
    if (op == OP_MUL && (is_number(left, 0.0) || is_number(right, 0.0))) {
        return become_number(e, 0.0);
    }
    double value;
    if (left->kind == EXPR_NUMBER && right->kind == EXPR_NUMBER
        && expr_apply(op, left->as.number, right->as.number, &value)) {
        return become_number(e, value);
    }
    return e;
}

Expr *expr_derive(const Expr *e, const char *name) {
    switch (e->kind) {
        case EXPR_NUMBER:
            return expr_number(0.0);
        case EXPR_VARIABLE:
            return expr_number(strcmp(e->as.name, name) == 0 ? 1.0 : 0.0);
        case EXPR_BINARY:
            break;
        case EXPR_IF:
        case EXPR_CALL:
            return NULL;
    }
    const Expr *u = e->as.binary.left;
    const Expr *v = e->as.binary.right;
    switch (e->as.binary.op) {
        case OP_ADD:
        case OP_SUB:
            return expr_binary(e->as.binary.op, expr_derive(u, name), expr_derive(v, name));
        case OP_MUL:
            return expr_binary(OP_ADD,
                               expr_binary(OP_MUL, expr_derive(u, name), expr_copy(v)),
                               expr_binary(OP_MUL, expr_copy(u), expr_derive(v, name)));
        case OP_DIV:
            return expr_binary(OP_DIV,
                               expr_binary(OP_SUB,
                                           expr_binary(OP_MUL, expr_derive(u, name), expr_copy(v)),
                                           expr_binary(OP_MUL, expr_copy(u), expr_derive(v, name))),
                               expr_binary(OP_MUL, expr_copy(v), expr_copy(v)));
        case OP_LESS:
        case OP_LESS_EQUAL:
        case OP_GREATER:
        case OP_GREATER_EQUAL:
        case OP_EQUAL:
        case OP_NOT_EQUAL:
            return NULL;
    }
    return NULL;
}