#include "fn.h"
#include "div.h"
#include "errloc.h"
#include "gl.h"
#include "im.h"
#include "parse.h"
#include "prv.h"
#include "pr.h"
#include "stmt.h"
#include "storage.h"
#include "td.h"
#include "ws.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[] =
{
    {(struct step[])
        {{ID, "public"}, {ID, "function"}, {ID}, {0}}, .fn = pubfn, .end = end},
    {(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, *stk = p->stk;

        if (lex_eof(l, p->stk))
            errloc(--stk, "function \"%s\" has no body", tk->s);
        else
            errloc(p->tk, "incomplete \"%s\"", p->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 tk *tk, const struct lex *lex, struct prv *p,
    enum linkage l)
{
    const struct tk *def;
    struct ast *ast = p->ast;
    size_t n = ast->nfns + 1;
    struct fn *fns;

    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;
    }

    fns[ast->nfns++] = (struct fn){.linkage = l, .tk = tk};
    ast->fns = fns;

    if (div_fn(lex, p))
        return -1;

    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(p->stk + 2, l, p, EXTERNAL);
}

static int prvfn(const struct lex *l, struct prv *p)
{
    return pushfn(p->stk + 1, l, p, STATIC);
}

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

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

    for (size_t i = 0; i < fn->nimps; i++)
        if ((e = fn_var(fn->imps[i].ast.fns, tk)))
            return e;

    return NULL;
}

int fn_cgen(const struct fn *fn, struct cgen *c)
{
    const struct pr *pr = fn->pr;
    const struct param *ret = pr->ret;
    const struct type *rt = NULL;
    const char *s = NULL;

    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)
    {
        rt = ret->entry->t;
        s = ret->entry->tk->s;
        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(" %%__param_%s ", p->tk->s);

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

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

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

        printf("%%%s_ =l alloc%zu %zu\n", s, t->align, t->sz);
        printf("store%s %%__param_%s, %%%s_\n", cgen_sz(t->sz), s, s);
    }

    if (ret)
        printf("%%%s_ =l alloc%zu %zu\n", s, rt->align, rt->sz);

    if (fn->ws && ws_cgen(fn->ws, c))
        return -1;

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

    if (ret)
    {
        const char *t = cgen_tmp(rt->sz);

        printf("%%%s =%s load%s %%%s_\n", t, cgen_sz(rt->sz), cgen_load(rt), s);
        printf("ret %%%s", t);
    }
    else
        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->nimps; i++)
        im_free(&fn->imps[i]);

    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->imps);
    free(fn->stmts);
}

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

    if (p->proto)
        return p->proto;

    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 (lex_eof(l, p->stk))
        return 0;
    else if (lex_eof(l, p->stk + 1))
    {
        errloc(p->stk, "incomplete function statement");
        return -1;
    }
    else if (push(&init, p))
        return -1;

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