#include 
#include "Python.h"
#include "Python-ast.h"

static PyObject *_str_open_br;
static PyObject *_str_dbl_open_br;
static PyObject *_str_close_br;
static PyObject *_str_dbl_close_br;

/* Forward declarations for recursion via helper functions. */
static PyObject *
expr_as_unicode(expr_ty e, int level);
static int
append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level);
static int
append_joinedstr(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec);
static int
append_formattedvalue(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec);
static int
append_ast_slice(_PyUnicodeWriter *writer, slice_ty slice);

static int
append_charp(_PyUnicodeWriter *writer, const char *charp)
{
    return _PyUnicodeWriter_WriteASCIIString(writer, charp, -1);
}

#define APPEND_STR_FINISH(str)  do { \
        return append_charp(writer, (str)); \
    } while (0)

#define APPEND_STR(str)  do { \
        if (-1 == append_charp(writer, (str))) { \
            return -1; \
        } \
    } while (0)

#define APPEND_STR_IF(cond, str)  do { \
        if ((cond) && -1 == append_charp(writer, (str))) { \
            return -1; \
        } \
    } while (0)

#define APPEND_STR_IF_NOT_FIRST(str)  do { \
        APPEND_STR_IF(!first, (str)); \
        first = false; \
    } while (0)

#define APPEND_EXPR(expr, pr)  do { \
        if (-1 == append_ast_expr(writer, (expr), (pr))) { \
            return -1; \
        } \
    } while (0)

#define APPEND(type, value)  do { \
        if (-1 == append_ast_ ## type(writer, (value))) { \
            return -1; \
        } \
    } while (0)

static int
append_repr(_PyUnicodeWriter *writer, PyObject *obj)
{
    int ret;
    PyObject *repr;
    repr = PyObject_Repr(obj);
    if (!repr) {
        return -1;
    }
    ret = _PyUnicodeWriter_WriteStr(writer, repr);
    Py_DECREF(repr);
    return ret;
}

/* Priority levels */

enum {
    PR_TUPLE,
    PR_TEST,            /* 'if'-'else', 'lambda' */
    PR_OR,              /* 'or' */
    PR_AND,             /* 'and' */
    PR_NOT,             /* 'not' */
    PR_CMP,             /* '', '==', '>=', 'v.Attribute.value->kind == Num_kind ||
        e->v.Attribute.value->kind == Constant_kind)
    {
        period = " .";
    }
    else {
        period = ".";
    }
    APPEND_STR(period);

    return _PyUnicodeWriter_WriteStr(writer, e->v.Attribute.attr);
}

static int
append_ast_simple_slice(_PyUnicodeWriter *writer, slice_ty slice)
{
    if (slice->v.Slice.lower) {
        APPEND_EXPR(slice->v.Slice.lower, PR_TEST);
    }

    APPEND_STR(":");

    if (slice->v.Slice.upper) {
        APPEND_EXPR(slice->v.Slice.upper, PR_TEST);
    }

    if (slice->v.Slice.step) {
        APPEND_STR(":");
        APPEND_EXPR(slice->v.Slice.step, PR_TEST);
    }
    return 0;
}

static int
append_ast_ext_slice(_PyUnicodeWriter *writer, slice_ty slice)
{
    Py_ssize_t i, dims_count;
    dims_count = asdl_seq_LEN(slice->v.ExtSlice.dims);
    for (i = 0; i < dims_count; i++) {
        APPEND_STR_IF(i > 0, ", ");
        APPEND(slice, (slice_ty)asdl_seq_GET(slice->v.ExtSlice.dims, i));
    }
    APPEND_STR_IF(dims_count == 1, ",");
    return 0;
}

static int
append_ast_index_slice(_PyUnicodeWriter *writer, slice_ty slice)
{
    int level = PR_TUPLE;
    expr_ty value = slice->v.Index.value;
    if (value->kind == Tuple_kind) {
        for (Py_ssize_t i = 0; i < asdl_seq_LEN(value->v.Tuple.elts); i++) {
            expr_ty element = asdl_seq_GET(value->v.Tuple.elts, i);
            if (element->kind == Starred_kind) {
                ++level;
                break;
            }
        }
    }
    APPEND_EXPR(value, level);
    return 0;
}

static int
append_ast_slice(_PyUnicodeWriter *writer, slice_ty slice)
{
    switch (slice->kind) {
    case Slice_kind:
        return append_ast_simple_slice(writer, slice);
    case ExtSlice_kind:
        return append_ast_ext_slice(writer, slice);
    case Index_kind:
        return append_ast_index_slice(writer, slice);
    default:
        PyErr_SetString(PyExc_SystemError,
                        "unexpected slice kind");
        return -1;
    }
}

static int
append_ast_subscript(_PyUnicodeWriter *writer, expr_ty e)
{
    APPEND_EXPR(e->v.Subscript.value, PR_ATOM);
    APPEND_STR("[");
    APPEND(slice, e->v.Subscript.slice);
    APPEND_STR_FINISH("]");
}

static int
append_ast_starred(_PyUnicodeWriter *writer, expr_ty e)
{
    APPEND_STR("*");
    APPEND_EXPR(e->v.Starred.value, PR_EXPR);
    return 0;
}

static int
append_ast_yield(_PyUnicodeWriter *writer, expr_ty e)
{
    if (!e->v.Yield.value) {
        APPEND_STR_FINISH("(yield)");
    }

    APPEND_STR("(yield ");
    APPEND_EXPR(e->v.Yield.value, PR_TEST);
    APPEND_STR_FINISH(")");
}

static int
append_ast_yield_from(_PyUnicodeWriter *writer, expr_ty e)
{
    APPEND_STR("(yield from ");
    APPEND_EXPR(e->v.YieldFrom.value, PR_TEST);
    APPEND_STR_FINISH(")");
}

static int
append_ast_await(_PyUnicodeWriter *writer, expr_ty e, int level)
{
    APPEND_STR_IF(level > PR_AWAIT, "(");
    APPEND_STR("await ");
    APPEND_EXPR(e->v.Await.value, PR_ATOM);
    APPEND_STR_IF(level > PR_AWAIT, ")");
    return 0;
}

static int
append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level)
{
    switch (e->kind) {
    case BoolOp_kind:
        return append_ast_boolop(writer, e, level);
    case BinOp_kind:
        return append_ast_binop(writer, e, level);
    case UnaryOp_kind:
        return append_ast_unaryop(writer, e, level);
    case Lambda_kind:
        return append_ast_lambda(writer, e, level);
    case IfExp_kind:
        return append_ast_ifexp(writer, e, level);
    case Dict_kind:
        return append_ast_dict(writer, e);
    case Set_kind:
        return append_ast_set(writer, e);
    case GeneratorExp_kind:
        return append_ast_genexp(writer, e);
    case ListComp_kind:
        return append_ast_listcomp(writer, e);
    case SetComp_kind:
        return append_ast_setcomp(writer, e);
    case DictComp_kind:
        return append_ast_dictcomp(writer, e);
    case Yield_kind:
        return append_ast_yield(writer, e);
    case YieldFrom_kind:
        return append_ast_yield_from(writer, e);
    case Await_kind:
        return append_ast_await(writer, e, level);
    case Compare_kind:
        return append_ast_compare(writer, e, level);
    case Call_kind:
        return append_ast_call(writer, e);
    case Constant_kind:
        return append_repr(writer, e->v.Constant.value);
    case Num_kind:
        return append_repr(writer, e->v.Num.n);
    case Str_kind:
        return append_repr(writer, e->v.Str.s);
    case JoinedStr_kind:
        return append_joinedstr(writer, e, false);
    case FormattedValue_kind:
        return append_formattedvalue(writer, e, false);
    case Bytes_kind:
        return append_repr(writer, e->v.Bytes.s);
    case Ellipsis_kind:
        APPEND_STR_FINISH("...");
    case NameConstant_kind:
        return append_repr(writer, e->v.NameConstant.value);
    /* The following exprs can be assignment targets. */
    case Attribute_kind:
        return append_ast_attribute(writer, e);
    case Subscript_kind:
        return append_ast_subscript(writer, e);
    case Starred_kind:
        return append_ast_starred(writer, e);
    case Name_kind:
        return _PyUnicodeWriter_WriteStr(writer, e->v.Name.id);
    case List_kind:
        return append_ast_list(writer, e);
    case Tuple_kind:
        return append_ast_tuple(writer, e, level);
    default:
        PyErr_SetString(PyExc_SystemError,
                        "unknown expression kind");
        return -1;
    }
}

static int
maybe_init_static_strings(void)
{
    if (!_str_open_br &&
        !(_str_open_br = PyUnicode_InternFromString("{"))) {
        return -1;
    }
    if (!_str_dbl_open_br &&
        !(_str_dbl_open_br = PyUnicode_InternFromString("{{"))) {
        return -1;
    }
    if (!_str_close_br &&
        !(_str_close_br = PyUnicode_InternFromString("}"))) {
        return -1;
    }
    if (!_str_dbl_close_br &&
        !(_str_dbl_close_br = PyUnicode_InternFromString("}}"))) {
        return -1;
    }
    return 0;
}

static PyObject *
expr_as_unicode(expr_ty e, int level)
{
    _PyUnicodeWriter writer;
    _PyUnicodeWriter_Init(&writer);
    writer.min_length = 256;
    writer.overallocate = 1;
    if (-1 == maybe_init_static_strings() ||
        -1 == append_ast_expr(&writer, e, level))
    {
        _PyUnicodeWriter_Dealloc(&writer);
        return NULL;
    }
    return _PyUnicodeWriter_Finish(&writer);
}

PyObject *
_PyAST_ExprAsUnicode(expr_ty e)
{
    return expr_as_unicode(e, PR_TEST);
}