#include "fn.h"
#include "div.h"
#include "errloc.h"
#include "gl.h"
#include "parse.h"
#include "prv.h"
#include "pr.h"
#include "stmt.h"
#include "storage.h"
#include "td.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

static int pubfn(const struct lex *l, struct prv *p);
static int prvfn(const struct lex *l, struct prv *p);
static int end(const struct lex *l, struct prv *p);

static const struct seq fnseq[] =
{
    {(const struct step[])
        {{ID, "function"}, {ID}, {ID, "public"}, {0}}, .fn = pubfn, .end = end},
    {(const struct step[])
        {{ID, "function"}, {ID}, {0}}, .fn = prvfn, .end = end},
    {0}
};

static int end(const struct lex *l, struct prv *p)
{
    const struct fn *fn = fn_cur(p);

    if (!fn->pr)
    {
        const struct tk *tk = fn->tk;

        errloc(tk, "function \"%s\" has no body", tk->s);
        return -1;
    }

    return 0;
}

static const struct tk *findfn(const char *s, const struct ast *ast)
{
    for (size_t i = 0; i < ast->nfns; i++)
    {
        const struct tk *tk = ast->fns[i].tk;

        if (!strcmp(tk->s, s))
            return tk;
    }

    return NULL;
}

static int pushfn(const struct lex *lex, struct prv *p, enum linkage l)
{
    const struct tk *tk = p->stk + 1, *def;
    struct ast *ast = p->ast;
    size_t n = ast->nfns + 1;
    struct fn *fns;
    struct pos init =
    {
        .seq = divseq,
        .stseq = divseq,
        .step = divseq->steps
    };

    if ((def = findfn(tk->s, ast)))
    {
        const struct loc *loc = &def->loc;

        errloc(tk, "function \"%s\" already defined at %s:%d:%d", tk->s,
            loc->f, loc->line, loc->col);
        return -1;
    }
    else if (!(fns = realloc(ast->fns, n * sizeof *fns)))
    {
        perror("realloc(3)");
        return -1;
    }
    else if (push(&init, p))
        return -1;

    fns[ast->nfns++] = (struct fn){.linkage = l, .tk = tk};
    ast->fns = fns;
    p->stk = p->tk;
    fprintf(stderr, "adding %s function %s\n",
        l == STATIC ? "private" : "public", tk->s);
    return 1;
}

static int pubfn(const struct lex *l, struct prv *p)
{
    return pushfn(l, p, EXTERNAL);
}

static int prvfn(const struct lex *l, struct prv *p)
{
    return pushfn(l, p, STATIC);
}

const struct stentry *fn_var(const struct fn *fn, const struct tk *tk)
{
    const struct stentry *e = NULL;
    const struct storage *lk = fn->lk, *ws = fn->ws, *gl = fn->gl;

    if ((!lk || !(e = storage_find(tk->s, lk)))
        && (!ws || !(e = storage_find(tk->s, ws)))
        && (!gl || !(e = storage_find(tk->s, gl))))
        return NULL;

    return e;
}

int fn_cgen(const struct fn *fn, struct cgen *c)
{
    const struct pr *pr = fn->pr;
    const struct param *ret = pr->ret;

    if (fn->gl && gl_cgen(fn->gl, c))
        return -1;

    switch (fn->linkage)
    {
    case STATIC:
        printf("function ");
        break;
    case EXTERNAL:
        printf("export function ");
        break;
    }

    if (ret)
        printf("%s ", cgen_abity(ret->entry->t));

    printf("$%s(", fn->tk->s);

    for (size_t i = 0; i < pr->nparams; i++)
    {
        const struct param *p = &pr->params[i];

        printf(" %s", cgen_abity(p->entry->t));
        printf(" %%%s_ ", p->tk->s);

        if (i + 1 < pr->nparams)
            fputs(", ", stdout);
    }

    puts(") {\n@start");

    for (size_t i = 0; i < fn->nstmts; i++)
        if (stmt_cgen(&fn->stmts[i], c))
            return -1;

    fputs("ret", stdout);

    if (pr->ret)
        fputs(" %__ret", stdout);

    puts("\n}");
    return 0;
}

void fn_cgen_free(struct fn_cgen *c)
{
    free(c);
}

void fn_free(struct fn *fn)
{
    for (size_t i = 0; i < fn->nstmts; i++)
        stmt_free(&fn->stmts[i]);

    storage_free(fn->lk);
    storage_free(fn->ws);
    storage_free(fn->gl);
    pr_free(fn->pr);
    td_free(fn->td);
    free(fn->stmts);
}

struct fn *fn_cur(const struct prv *p)
{
    struct ast *ast = p->ast;

    return &ast->fns[p->ast->nfns - 1];
}

int fn(const struct lex *l, struct prv *p)
{
    struct pos init =
    {
        .seq = fnseq,
        .stseq = fnseq,
        .step = fnseq->steps
    };

    if (push(&init, p))
        return -1;

    p->stk = p->tk;
    return 1;
}