#include <math.h>
#include <limits.h>
#include <stdio.h>
#include "util/err.h"
#include "parse.h"
#include "lex.h"

static bool is_expr(struct token *tok)
{
    return tok->type == TOKEN_NUMBER || tok->type == TOKEN_VAR || (tok->type == TOKEN_OP && op_category_tab[tok->op.type] == OPC_EXPR);
}

static struct var *get_var(struct token *tok, struct source *src, struct registry *reg)
{
    struct var *var = nullptr;
    map_find(&reg->vars, tok->var, &var);
    if (!var)
        source_error(src, tok->loc, "no such variable: $%.*s", PSTR(tok->var));    
    return var;
}

static void parse_expr(struct token *tok, struct expr *expr, struct source *src, struct registry *reg, bool has_time)
{
    expr->loc = tok->loc;
    if (tok->type == TOKEN_NUMBER) {
        expr->type = EXPR_NUMBER;
        expr->number = tok->number;
    } else if (tok->type == TOKEN_VAR) {
        expr->type = EXPR_NUMBER;
        struct var *var = get_var(tok, src, reg);
        if (var->type != VAR_NUMBER)
            source_error(src, tok->loc, "$%.*s is not a number", PSTR(tok->var));
        expr->number = var->number;
    } else if (tok->type == TOKEN_OP) {
        size_t n_args = tok->op.children.len;
        if (tok->op.type == OP_TIME) {
            if (!has_time)
                source_error(src, tok->loc, "'time' not allowed in this context");
            if (n_args)
                source_error(src, tok->loc, "'time' does not expect any arguments");
            expr->type = EXPR_TIME;
        } else {
            switch (tok->op.type) {
                case OP_ADD:
                case OP_MUL:
                    break;
                case OP_MIN:
                case OP_MAX:
                    if (n_args < 1)
                        source_error(src, tok->loc, "'%.*s' expects at least 1 argument",
                            PSTR(op_name_tab[tok->op.type]));
                    break;
                case OP_DIV:
                case OP_SUB:
                    if (n_args != 1 && n_args != 2)
                        source_error(src, tok->loc, "'%.*s' expects one or two arguments",
                            PSTR(op_name_tab[tok->op.type]));
                    break;
                case OP_MOD:
                    if (n_args != 1)
                        source_error(src, tok->loc, "'%.*s' expects exactly two arguments",
                            PSTR(op_name_tab[tok->op.type]));
                    break;
                case OP_CLAMP:
                case OP_MIX:
                    if (n_args != 3)
                        source_error(src, tok->loc, "'%.*s' expects exactly three arguments",
                            PSTR(op_name_tab[tok->op.type]));
                    break;
                case OP_ABS:
                    if (n_args != 1)
                        source_error(src, tok->loc, "'%.*s' expects exactly one argument",
                            PSTR(op_name_tab[tok->op.type]));
                    break;
                default:
                    source_error(src, tok->op.name_loc, "not a valid math operation");
            }
            expr->type = EXPR_OP;
            expr->op.type = tok->op.type;
            array_alloc(&expr->op.args, n_args);
            for (size_t i = 0; i < n_args; i++)
                parse_expr(&tok->op.children.ptr[i], &expr->op.args.ptr[i], src, reg, has_time);
        }
    } else {
        source_error(src, tok->loc, "not a valid math expression");
    }
}

static double parse_eval_expr(struct token *tok, struct source *src, struct registry *reg)
{
    struct expr expr;
    parse_expr(tok, &expr, src, reg, false);
    double num = expr_eval(&expr, 0.0);
    expr_free(&expr);
    return num;
}

double parse_eval_expr_range(struct token *tok, struct source *src, struct registry *reg, double low, double high)
{
    double num = parse_eval_expr(tok, src, reg);
    if (num < low)
        source_error(src, tok->loc, "%f is smaller than %f", num, low);
    if (num > high)
        source_error(src, tok->loc, "%f is larger than %f", num, high);
    return num;    
}

unsigned int parse_eval_expr_uint(struct token *tok, struct source *src, struct registry *reg)
{
    double num = round(parse_eval_expr(tok, src, reg));
    if (num < 0.0 || num > UINT_MAX)
        source_error(src, tok->loc, "%f cannot be converted to an unsigned integer", num);
    return (unsigned int) num;
}

unsigned int parse_eval_expr_uint_range(struct token *tok, struct source *src, struct registry *reg, unsigned int low, unsigned int high)
{
    unsigned int num = parse_eval_expr_uint(tok, src, reg);
    if (low && num < low)
        source_error(src, tok->loc, "%u is smaller than %u", num, low);
    if (high && num > high)
        source_error(src, tok->loc, "%u is larger than %u", num, high);
    return num;
}

void parse_string(struct token *tok, str *out, struct source *src, struct registry *reg)
{
    if (tok->type == TOKEN_STRING) {
        *out = tok->string;
    } else if (tok->type == TOKEN_VAR) {
        struct var *var = get_var(tok, src, reg);
        if (var->type != VAR_STRING)
            source_error(src, tok->loc, "$%.*s is not a string", PSTR(tok->var));
        *out = var->string;
    } else {
        source_error(src, tok->loc, "not a string");
    }
}

void parse_string_loc(struct token *tok, struct loc_str *out, struct source *src, struct registry *reg)
{
    out->loc = tok->loc;
    parse_string(tok, &out->str, src, reg);
}

static void parse_name(struct token *tok, struct loc_str *name, struct source *src)
{
    if (tok->type != TOKEN_IDENT)
        source_error(src, tok->loc, "not an identifier");
    name->loc = tok->loc;
    name->str = tok->ident;
}

void parse_seq_inst(struct token *tok, struct seq_inst *inst, struct source *src, struct registry *reg)
{
    if (is_expr(tok)) {
        inst->type = SEQ_TIME;
        inst->time = parse_eval_expr_range(tok, src, reg, 0.0, INFINITY);
    } else if (tok->type == TOKEN_OP) {
        if (tok->op.type == OP_POV) {
            if (tok->op.children.len != 1)
                source_error(src, tok->loc, "'pov' expects exactly one argument");
            inst->type = SEQ_POV;
            parse_name(&tok->op.children.ptr[0], &inst->pov_name, src);
        } else if (tok->op.type == OP_REP) {
            if (tok->op.children.len < 2)
                source_error(src, tok->loc, "'rep' expects at least two arguments");

            inst->type = SEQ_REP;

            struct token *tok_times = &tok->op.children.ptr[0];
            if (is_expr(tok_times))
                inst->rep.times = parse_eval_expr_uint(tok_times, src, reg);
            else if (tok_times->type == TOKEN_OP && tok_times->op.type == OP_INF)
                inst->rep.times = SEQ_TIMES_INF;
            else
                source_error(src, tok_times->loc, "not a valid number of repetitions");

            size_t n_child = tok->op.children.len-1;
            array_alloc(&inst->rep.children, n_child);
            for (size_t i = 0; i < n_child; i++)
                parse_seq_inst(&tok->op.children.ptr[i+1], &inst->rep.children.ptr[i], src, reg);
        } else {
            source_error(src, tok->op.name_loc, "not a valid seq instruction");
        }
    } else {
        source_error(src, tok->loc, "not a valid seq instruction");
    }
}

static bool is_transform_op(struct token *tok)
{
    return tok->type == TOKEN_OP && op_category_tab[tok->op.type] == OPC_TRANSFORM;
}

static char swizzle_char[3] = { 'x', 'y', 'z' };
static str transform_op_name_tab[] = { S("translate"), S("scale"), S("rotate") };

static void transform_parse_swizzle(size_t n, uint8_t swizzle[3], struct token *tok, struct source *src)
{
    char c = tok->ident.ptr[n];
    for (size_t d = 0; d < 3; d++) {
        if (c != swizzle_char[d])
            continue;
        for (size_t i = 0; i < n; i++)
            if (swizzle[i] == d)
                source_error(src, tok->loc, "'%c' occurs twice in swizzle", c);
        swizzle[n] = d;
        return;
    }
    source_error(src, tok->loc, "invalid swizzle component: '%c'", c);
}

static void transform_check_assign(struct transform *trans, enum transform_op op, uint8_t dim, struct token *tok, struct source *src)
{
    if (trans->has_op[op][dim])
        source_error(src, tok->loc, "%.*s %c already assigned",
            PSTR(transform_op_name_tab[op]), swizzle_char[dim]);
    trans->has_op[op][dim] = true;
}

static void parse_transform_op(struct token *tok, struct transform *trans, struct source *src, struct registry *reg)
{
    if (tok->type != TOKEN_OP)
        source_error(src, tok->loc, "not a valid transform operator");
    enum transform_op op;
    switch (tok->op.type) {
        case OP_TRANSLATE:
            op = TR_TRANSLATE;
            break;
        case OP_SCALE:
            op = TR_SCALE;
            break;
        case OP_ROTATE:
            op = TR_ROTATE;
            break;
        default:
            source_error(src, tok->op.name_loc, "not a valid transform operator");
    }

    if (tok->op.children.len < 1)
        source_error(src, tok->op.name_loc, "'%.*s' expects at least one argument",
            PSTR(op_name_tab[tok->op.type]));

    size_t offset = 0;
    uint8_t swizzle[3];
    size_t n_swizzle;

    struct token *tok_swizzle = &tok->op.children.ptr[0];
    if (tok_swizzle->type == TOKEN_IDENT) {
        n_swizzle = tok_swizzle->ident.len;
        if (n_swizzle > 3)
            source_error(src, tok_swizzle->loc, "swizzle must not be longer than 3 chars");
        for (size_t i = 0; i < n_swizzle; i++)
            transform_parse_swizzle(i, swizzle, tok_swizzle, src);
        offset = 1;

        if (tok->op.children.len-offset != n_swizzle) {
            if (n_swizzle == 1)
                source_error(src, tok->loc, "expected exactly one numeric argument after swizzle");
            else if (tok->op.children.len-offset != 1)
                source_error(src, tok->loc, "expected 1 or %zu numeric arguments after swizzle", n_swizzle);
        }
    } else {
        n_swizzle = 3;
        for (uint8_t i = 0; i < n_swizzle; i++)
            swizzle[i] = i;

        if (tok->op.children.len != 1 && tok->op.children.len != 3)
            source_error(src, tok->loc, "'%.*s' without a swizzle expects either one or three arguments",
                PSTR(op_name_tab[tok->op.type]));
    }

    for (size_t i = offset; i < tok->op.children.len; i++) {
        uint8_t dim = swizzle[i-offset];
        transform_check_assign(trans, op, dim, tok, src);
        parse_expr(&tok->op.children.ptr[i], &trans->op[op][dim], src, reg, true);
    }

    if (n_swizzle > 1 && tok->op.children.len-offset == 1) {
        for (uint8_t i = 1; i < n_swizzle; i++) {
            transform_check_assign(trans, op, swizzle[i], tok, src);
            expr_copy(&trans->op[op][swizzle[i]], &trans->op[op][swizzle[0]]);
        }
    }
}

/*
static bool is_pos(struct token *tok)
{
    return tok->type == TOKEN_OP && op_category_tab[tok->op.type] == OPC_POS;
}
*/

static void parse_pos(struct token *tok, struct pos *pos, struct source *src, struct registry *reg)
{
    if (tok->type != TOKEN_OP)
        source_error(src, tok->loc, "not a valid position");
    switch (tok->op.type) {
        case OP_POS:
            if (tok->op.children.len != 3)
                source_error(src, tok->loc, "'pos' expects exactly three arguments");          
            pos->type = POS_ABSOLUTE;
            for (size_t i = 0; i < tok->op.children.len; i++)
                parse_expr(&tok->op.children.ptr[i], &pos->absolute[i], src, reg, true);
            break;
        case OP_ATTACH:
            if (tok->op.children.len != 1)
                source_error(src, tok->loc, "'attach' expects exactly one argument");
            pos->type = POS_ATTACHED;
            parse_name(&tok->op.children.ptr[0], &pos->attach_name, src);
            break;
        default:
            source_error(src, tok->op.name_loc, "not a valid position");
    }
}

static void parse_surface(size_t n_attrs, struct token attrs[n_attrs], struct surface *surf, struct source *src, struct registry *reg)
{
    surface_init(surf);
    for (size_t i = 0; i < n_attrs; i++) {
        struct token *tok = &attrs[i];
        if (tok->type != TOKEN_OP)
            source_error(src, tok->loc, "not a valid surface attribute");
        switch (tok->op.type) {
            case OP_TEXTURE:
            case OP_CUBEMAP:
                if (tok->op.children.len != 1)
                    source_error(src, tok->loc, "'%.*s' expects exactly one argument",
                        PSTR(op_name_tab[tok->op.type]));
                if (surf->texture.present)
                    source_error(src, tok->loc, "texture already set");
                surf->texture.present = true;
                surf->texture.data.source.type = tok->op.type == OP_CUBEMAP ? TEXTURE_CUBEMAP : TEXTURE_2D;
                parse_string_loc(&tok->op.children.ptr[0], &surf->texture.data.source.path, src, reg);
                break;
            case OP_COLOR:
                if (tok->op.children.len != 1)
                    source_error(src, tok->loc, "'color' expects exactly one argument");   
                expr_free(&surf->color);
                parse_expr(&tok->op.children.ptr[0], &surf->color, src, reg, true);
                break;
            default:
                source_error(src, tok->op.name_loc, "not a valid surface attribute");            
        }
    }
}

static void parse_node(struct token *tok, struct node *node, struct source *src, struct registry *reg)
{
    if (tok->type != TOKEN_OP)
        source_error(src, tok->loc, "not a valid scene node");
    switch (tok->op.type) {
        case OP_TRANSFORM: {
            node->type = NODE_TRANSFORM;
            memset(&node->transform.has_op, 0, sizeof node->transform.has_op);

            size_t idx = 0;
            for (; idx < tok->op.children.len && is_transform_op(&tok->op.children.ptr[idx]); idx++)
                parse_transform_op(&tok->op.children.ptr[idx], &node->transform, src, reg);

            array_alloc(&node->transform.children, tok->op.children.len-idx);
            for (size_t i = idx; i < tok->op.children.len; i++)
                parse_node(&tok->op.children.ptr[i], &node->transform.children.ptr[i-idx], src, reg);
            break;
        }
        case OP_CAMERA: {
            if (tok->op.children.len < 1)
                source_error(src, tok->loc, "'camera' expects at least one argument");
            size_t idx = 0;
            struct token *tok_name = &tok->op.children.ptr[idx];
            if (tok_name->type == TOKEN_IDENT)
                idx++;
            else
                tok_name = nullptr;
            if (tok->op.children.len < idx+1)
                source_error(src, tok->loc, "named camera expects at least two arguments");
            node->type = NODE_CAMERA;
            parse_pos(&tok->op.children.ptr[idx], &node->camera.target, src, reg);
            node->camera.roll = false;
            if (tok->op.children.len == idx+2) {
                struct token *tok_attr = &tok->op.children.ptr[idx+1];
                if (tok_attr->type != TOKEN_OP)
                    source_error(src, tok_attr->loc, "not a valid camera attribute");
                if (tok_attr->op.type != OP_ROLL)
                    source_error(src, tok_attr->op.name_loc, "not a valid camera attribute");
                if (tok_attr->op.children.len != 0)
                    source_error(src, tok_attr->loc, "'roll' does not expect any arguments");
                node->camera.roll = true;
            } else if (tok->op.children.len != idx+1) {
                source_error(src, tok->loc, "excess camera arguments");
            }
            if (tok_name) {
                bool succ;
                map_insert(&reg->cameras, tok_name->ident, &node->camera, &succ);
                if (!succ)
                    source_error(src, tok_name->loc, "camera '%.*s' already registered",
                        PSTR(tok_name->ident));
            } else {
                if (reg->default_camera)
                    source_error(src, tok->loc, "unnamed camera already registered");
                reg->default_camera = &node->camera;
            }
            break;
        }
        case OP_CUBE:
            node->type = NODE_OBJECT;
            node->object.mesh.source.type = MESH_CUBE;
            parse_surface(tok->op.children.len, tok->op.children.ptr, &node->object.surface, src, reg);
            break;
        case OP_OBJ:
            if (tok->op.children.len < 1)
                source_error(src, tok->loc, "'obj' expects at least one argument");
            node->type = NODE_OBJECT;
            node->object.mesh.source.type = MESH_OBJ;
            parse_string_loc(&tok->op.children.ptr[0], &node->object.mesh.source.path, src, reg);
            parse_surface(tok->op.children.len-1, tok->op.children.ptr+1, &node->object.surface, src, reg);
            break;
        case OP_BONE: {
            node->type = NODE_BONE;
            if (tok->op.children.len != 1)
                source_error(src, tok->loc, "'bone' expects exactly one argument");
            struct token *tok_name = &tok->op.children.ptr[0];
            if (tok_name->type != TOKEN_IDENT)
                source_error(src, tok_name->loc, "not an identifier");
            bool succ;
            map_insert(&reg->bones, tok_name->ident, &node->bone, &succ);
            if (!succ)
                source_error(src, tok_name->loc, "bone '%.*s' already registered",
                    PSTR(tok_name->ident));
            break;
        }
        case OP_LIGHT:
            node->type = NODE_LIGHT;
            source_error(src, tok->loc, "TODO: implement light");
            break;
        default:
            source_error(src, tok->op.name_loc, "not a valid scene node");
    }
}

static void parse_toplevel(struct token *tok, struct scene *scene, struct source *src, struct registry *reg)
{
    if (tok->type != TOKEN_OP)
        source_error(src, tok->loc, "not allowed at top-level");

    switch (tok->op.type) {
        case OP_VAR: {
            if (tok->op.children.len != 2 && tok->op.children.len != 3)
                source_error(src, tok->loc, "'var' expects two or three arguments");

            struct var var;

            struct token *tok_type = &tok->op.children.ptr[0];
            if (tok_type->type == TOKEN_IDENT && array_eq(tok_type->ident, S("string")))
                var.type = VAR_STRING;
            else if (tok_type->type == TOKEN_IDENT && array_eq(tok_type->ident, S("number")))
                var.type = VAR_NUMBER;
            else
                source_error(src, tok_type->loc, "not a valid variable type");

            struct token *tok_name = &tok->op.children.ptr[1];
            if (tok_name->type != TOKEN_VAR)
                source_error(src, tok_name->loc, "not a variable name");

            if (tok->op.children.len == 3) {
                struct token *tok_default = &tok->op.children.ptr[2];
                switch (var.type) {
                    case VAR_STRING: parse_string(tok_default, &var.string, src, reg); break;
                    case VAR_NUMBER: var.number = parse_eval_expr(tok_default, src, reg); break;
                }
            }

            str *raw_var = nullptr;
            map_find(&reg->raw_vars, tok_name->var, &raw_var);

            if (raw_var) {
                switch (var.type) {
                    case VAR_STRING:
                        var.string = *raw_var;
                        break;
                    case VAR_NUMBER:
                        if (!str_parse_double(*raw_var, &var.number))
                            eprintf("invalid value for variable $%.*s: not a number: %.*s\n",
                                PSTR(tok_name->var), PSTR(*raw_var));
                        break;
                }
            } else if (tok->op.children.len != 3) {
                eprintf("missing value for %.*s variable $%.*s\n",
                    PSTR(tok_type->ident), PSTR(tok_name->var));
            }

            bool success;
            map_insert(&reg->vars, tok_name->var, var, &success);
            if (!success)
                source_error(src, tok->loc, "variable $%.*s already declared", PSTR(tok_name->var));
            break;
        }
        case OP_FPS:
            if (tok->op.children.len != 1)
                source_error(src, tok->loc, "'fps' expects exactly one argument");
            scene->fps = parse_eval_expr_uint_range(&tok->op.children.ptr[0], src, reg, 1, 0);
            break;
        case OP_SIZE:
            if (tok->op.children.len != 2)
                source_error(src, tok->loc, "'size' expects exactly two arguments");
            for (size_t i = 0; i < 2; i++)
                scene->size[i] = parse_eval_expr_uint_range(&tok->op.children.ptr[i], src, reg, 1, 0);
            break;
        case OP_SEQ:
            if (scene->seq.ptr)
                source_error(src, tok->loc, "sequence already defined");
            array_alloc(&scene->seq, tok->op.children.len);
            for (size_t i = 0; i < tok->op.children.len; i++)
                parse_seq_inst(&tok->op.children.ptr[i], &scene->seq.ptr[i], src, reg);
            break;
        case OP_PROJ: {
            if (tok->op.children.len < 1)
                source_error(src, tok->loc, "'proj' expects at least one argument");

            enum projection_type type;
            struct token *tok_type = &tok->op.children.ptr[0];
            if (tok_type->type == TOKEN_IDENT && array_eq(tok_type->ident, S("perspective")))
                type = PROJ_PERSPECTIVE;
            else if (tok_type->type == TOKEN_IDENT && array_eq(tok_type->ident, S("ortho")))
                type = PROJ_ORTHO;
            else
                source_error(src, tok_type->loc, "not a valid projection type");

            proj_free(&scene->proj);
            proj_default(&scene->proj, type);

            for (size_t i = 1; i < tok->op.children.len; i++) {
                struct token *child = &tok->op.children.ptr[i];
                if (child->type == TOKEN_OP && child->op.type == OP_FOV) {
                    if (type != PROJ_PERSPECTIVE)
                        source_error(src, child->op.name_loc, "'fov' only valid for perspective projection");
                    if (child->op.children.len != 1)
                        source_error(src, child->loc, "'fov' expects exactly one argument");
                    expr_free(&scene->proj.perspective.fov);
                    parse_expr(&child->op.children.ptr[0], &scene->proj.perspective.fov, src, reg, true);
                } else if (child->type == TOKEN_OP && child->op.type == OP_PLANES) {
                    float *planes;
                    switch (type) {
                        case PROJ_PERSPECTIVE:
                            if (child->op.children.len != 2)
                                source_error(src, child->loc, "'planes' expects exactly two arguments for perspective projection");
                            planes = scene->proj.perspective.planes;
                            break;
                        case PROJ_ORTHO:
                            if (child->op.children.len != 6)
                                source_error(src, child->loc, "'planes' expects exactly six arguments for ortho projection");
                            planes = scene->proj.ortho.planes;
                            break;
                    }
                    for (size_t i = 0; i < child->op.children.len; i++)
                        planes[i] = parse_eval_expr(&child->op.children.ptr[i], src, reg);
                } else {
                    source_error(src, child->loc, "not a valid projection attribute");
                }
            }
            break;
        }
        case OP_OUT:
            if (tok->op.children.len != 1)
                source_error(src, tok->loc, "'out' expects exactly one argument");
            free(scene->outfile.ptr);
            parse_string(&tok->op.children.ptr[0], &scene->outfile, src, reg);
            break;
        case OP_SCENE:
            if (scene->root.transform.children.ptr)
                source_error(src, tok->loc, "scene already defined");
            array_alloc(&scene->root.transform.children, tok->op.children.len);
            for (size_t i = 0; i < tok->op.children.len; i++)
                parse_node(&tok->op.children.ptr[i], &scene->root.transform.children.ptr[i], src, reg);
            break;
        default:
            source_error(src, tok->op.name_loc, "not allowed at top-level");
    }
}

void parse_scene(struct source *src, struct scene *scene, struct registry *reg)
{
    struct token tok;
    while (lex_token(src, &tok)) {
        parse_toplevel(&tok, scene, src, reg);
        free_token(&tok);
    }
}