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

static int pryy(const struct lex *l, struct prv *p);
static int pryn(const struct lex *l, struct prv *p);
static int prny(const struct lex *l, struct prv *p);
static int prnn(const struct lex *l, struct prv *p);

static const struct seq seqs[] =
{
    /* function definition with parameters and return variable */
    {
        (const struct step[])
        {
            {ID, "with"},
            {ID, NULL, 1},
            {ID, "returning"},
            {ID},
            {0}
        },
        .fn = pryy
    },

    /* function definition with parameters, no return variable */
    {
        (const struct step[])
        {
            {ID, "with"},
            {ID, NULL, 1},
            {0}
        },

        .fn = pryn
    },

    /* function definition with no parameters, with return variable */
    {
        (const struct step[])
        {
            {ID, "returning"},
            {ID},
            {0}
        },

        .fn = prny
    },

    {0}
};

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

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

    return NULL;
}

static int param(const struct tk *tk, const struct stentry *e, struct pr *pr)
{
    size_t n = pr->nparams + 1;
    struct param *params = realloc(pr->params, n * sizeof *params);

    if (!params)
    {
        perror("realloc(3)");
        return -1;
    }

    pr->params = params;
    pr->params[pr->nparams++] = (struct param){.tk = tk, .entry = e};
    return 0;
}

static int setret(const struct tk *tk, const struct stentry *e, struct pr *pr)
{
    if (!(pr->ret = malloc(sizeof *pr->ret)))
    {
        perror("malloc(3)");
        return -1;
    }

    *pr->ret = (struct param){.tk = tk, .entry = e};
    return 0;
}

static int finalize(struct prv *p)
{
    const struct fn *fn = fn_cur(p);
    const struct pr *pr = fn->pr;
    struct pos init =
    {
        .seq = stmts,
        .stseq = stmts,
        .step = stmts->steps
    };

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

    fprintf(stderr, "\tadding pr with %zu parameters and%s return value\n",
        pr->nparams, pr->ret ? "" : " no");
    p->stk = p->tk;
    return 1;
}

static int prret(const struct lex *l, const struct tk *tk, struct prv *p)
{
    const struct fn *fn = fn_cur(p);
    struct pr *pr = fn->pr;
    const char *s = tk->s;
    const struct storage *lk = fn->lk, *gl = fn->gl;
    const struct stentry *e, *gle;
    const struct tk *prm;

    if (!lk || !(e = storage_find(s, lk)))
    {
        errloc(tk, "return value \"%s\" not in %s", s, lk->name);
        return -1;
    }
    else if (gl && (gle = storage_find(s, gl)))
    {
        const struct loc *loc = &gle->tk->loc;

        errloc(tk, "return value \"%s\" also in %s at %d:%d", s, gl->name,
            loc->line, loc->col);
        return -1;
    }
    else if ((prm = findparam(s, pr)))
    {
        const struct loc *loc = &prm->loc;

        errloc(tk, "return value \"%s\" already listed as "
            "parameter at %s:%d:%d", tk->s, loc->f, loc->line, loc->col);
        return -1;
    }
    else if (setret(tk, e, pr))
        return -1;

    return finalize(p);
}

static const struct tk *prparams(const struct lex *l, struct prv *p, int *ret)
{
    const struct tk *tk = p->stk + 1;
    const struct fn *fn = fn_cur(p);
    struct pr *pr = fn->pr;

    if (ret)
        *ret = 0;

    for (;;)
    {
        const char *s = tk->s;
        const struct stentry *e, *gle;
        const struct tk *prm;
        const struct storage *lk = fn->lk, *gl = fn->gl;

        if (kw(s) && strcmp(s, "returning"))
            return tk;
        else if (!strcmp(s, "returning") && ret)
        {
            *ret = 1;
            return ++tk;
        }
        else if ((prm = findparam(s, pr)))
        {
            const struct loc *loc = &prm->loc;

            errloc(tk, "parameter \"%s\" already listed at %d:%d", s,
                loc->line, loc->col);
            return NULL;
        }
        else if (!lk || !(e = storage_find(s, lk)))
        {
            errloc(tk, "parameter \"%s\" not in %s", s, lk->name);
            return NULL;
        }
        else if (gl && (gle = storage_find(s, gl)))
        {
            const struct loc *loc = &gle->tk->loc;

            errloc(tk, "parameter \"%s\" also in %s at %d:%d", s, gl->name,
                loc->line, loc->col);
            return NULL;
        }
        else if (param(tk, e, pr))
            return NULL;
        else
            tk++;
    }

    return tk;
}

static int pryy(const struct lex *l, struct prv *p)
{
    int ret;
    const struct tk *tk = prparams(l, p, &ret);

    if (!tk)
        return -1;
    else if (ret)
        return prret(l, tk, p);

    return finalize(p);
}

static int pryn(const struct lex *l, struct prv *p)
{
    if (!prparams(l, p, NULL))
        return -1;

    return finalize(p);
}

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

static int prnn(const struct lex *l, struct prv *p)
{
    return finalize(p);
}

void pr_free(struct pr *pr)
{
    if (pr)
    {
        free(pr->params);
        free(pr->ret);
    }

    free(pr);
}

int pr(const struct lex *l, struct prv *p)
{
    const struct tk *tk = p->tk - 1, *next = p->tk;
    struct fn *fn = fn_cur(p);
    struct pos init =
    {
        .seq = seqs,
        .stseq = seqs,
        .step = seqs->steps
    };

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

        errloc(tk, "pr already defined at %s:%d:%d", loc->f, loc->line,
            loc->col);
        return -1;
    }
    else if (!((fn->pr = malloc(sizeof *fn->pr))))
    {
        perror("malloc(3)");
        return -1;
    }

    *fn->pr = (struct pr){.tk = tk};

    /* on no parameters or return values, go directly to next sequences list. */
    if (next - l->tokens < l->ntok
        && strcmp(next->s, "with")
        && strcmp(next->s, "returning"))
        return prnn(l, p);
    else if (push(&init, p))
        return -1;

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