#define PY_SSIZE_T_CLEAN
#include <Python.h>

#include "cmunipack.h"
#include "catfileclass.h"

PyObject *cmpack_error(int code);
PyObject *err_closed(void);
PyObject *return_datetime(CmpackDateTime *dt);

static PyObject *
CatFile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    CatFile *self = (CatFile*)type->tp_alloc(type, 0);
    if (self != NULL) 
		self->f = NULL;
    return (PyObject *)self;
}

static void
CatFile_dealloc(CatFile* self)
{
	if (self->f) {
		cmpack_cat_destroy(self->f);
		self->f = NULL;
	}
    self->ob_base.ob_type->tp_free((PyObject*)self);
}

static int
CatFile_init(CatFile *self, PyObject *args)
{
    if (!PyArg_ParseTuple(args, ""))
        return -1;

	if (!self->f) {
		self->f = cmpack_cat_new();
		if (!self->f) {
			PyErr_SetString(PyExc_MemoryError, "Memory allocation error");
			return -1;
		}
	}
    return 0;
}

static PyObject *
CatFile_open(CatFile *self, PyObject *args)
{
    char *name; 
	int res, mode;
	CmpackCatFile *f;

    if (!PyArg_ParseTuple(args, "esi", Py_FileSystemDefaultEncoding, &name, &mode))
        return NULL;

	res = cmpack_cat_open(&f, name, mode, 0);
	PyMem_Free(name);
	if (res!=0) 
		return cmpack_error(res);

	return CatFile_FromHandle(f);
}

static PyObject *
CatFile_test(CatFile *self, PyObject *args)
{
	char *filename;
	PyObject *res;

	if (!PyArg_ParseTuple(args, "es", Py_FileSystemDefaultEncoding, &filename))
		return NULL;

	res = PyBool_FromLong(cmpack_cat_test(filename)!=0);
	PyMem_Free(filename);
	return res;
}

static PyObject *
CatFile_self(CatFile *self)
{
    if (!self->f)
        return err_closed();
    Py_INCREF(self);
    return (PyObject *)self;
} 

static PyObject *
CatFile_exit(PyObject *f, PyObject *args)
{
    PyObject *ret = PyObject_CallMethod(f, "close", NULL);
    if (!ret)
        return NULL;

    Py_DECREF(ret);
    /* We cannot return the result of close since a true
     * value will be interpreted as "yes, swallow the
     * exception if one was raised inside the with block". */
    Py_RETURN_NONE;
} 

static PyObject *
CatFile_close(CatFile *self, PyObject *args)
{
	if (self->f) {
		int res = cmpack_cat_close(self->f);
		self->f = NULL;
		if (res!=0)
			return cmpack_error(res);
	}
	Py_RETURN_NONE;
}

static PyObject *
CatFile_copy(CatFile *self, PyObject *args)
{
	int res;
	CmpackCatFile *newf;

	if (!self->f)
		return err_closed();

	newf = cmpack_cat_new();
	if (!newf) {
		PyErr_SetString(PyExc_MemoryError, "Memory allocation error.");
		return NULL;
	}
	res = cmpack_cat_copy(newf, self->f);
	if (res!=0) {
		cmpack_cat_destroy(newf);
		return cmpack_error(res);
	}
	return CatFile_FromHandle(newf);
}

static PyObject *
CatFile_save(CatFile *self, PyObject *args, PyObject *kwds)
{
	int res;
    char *to; 
	CmpackCatFile *f;

    static char *kwlist[] = {"to", NULL};
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "es", kwlist, 
		Py_FileSystemDefaultEncoding, &to))
        return NULL;

	if (!self->f) {
		PyMem_Free(to);
		return err_closed();
	}

	res = cmpack_cat_open(&f, to, CMPACK_OPEN_CREATE, 0);
	if (res==0)
		res = cmpack_cat_copy(f, self->f);
	cmpack_cat_destroy(f);
	PyMem_Free(to);
	if (res!=0)
		return cmpack_error(res);
	Py_RETURN_NONE;
}

static PyObject *
CatFile_image_size(CatFile *self, void *data)
{
	if (!self->f)
		return err_closed();

	return Py_BuildValue("ii", cmpack_cat_get_width(self->f), 
		cmpack_cat_get_height(self->f));
}

static PyObject *
CatFile_date_time(CatFile *self, void *data)
{
	double jd;
	CmpackDateTime dt;

	if (!self->f)
		return err_closed();

	if (cmpack_cat_gkyd(self->f, "jd", &jd)==0) {
		if (cmpack_decodejd(jd, &dt)==0)
			return return_datetime(&dt);
	}
	Py_RETURN_NONE;
}

static PyObject *
CatFile_julian_date(CatFile *self, void *data)
{
	double jd;

	if (!self->f)
		return err_closed();

	if (cmpack_cat_gkyd(self->f, "jd", &jd)==0)
		return PyFloat_FromDouble(jd);
	Py_RETURN_NONE;
}

static PyObject *
CatFile_exposure(CatFile *self, void *data)
{
	double exptime;

	if (!self->f)
		return err_closed();

	if (cmpack_cat_gkyd(self->f, "exptime", &exptime)==0)
		return PyFloat_FromDouble(exptime);
	Py_RETURN_NONE;
}

static PyObject *
CatFile_filter(CatFile *self, void *data)
{
	const char *filter;

	if (!self->f)
		return err_closed();

	filter = cmpack_cat_gkys(self->f, "filter");
	if (filter)
		return PyUnicode_FromString(filter);
	Py_RETURN_NONE;
}

static PyObject *
CatFile_observer(CatFile *self, void *data)
{
	const char *filter;

	if (!self->f)
		return err_closed();

	filter = cmpack_cat_gkys(self->f, "observer");
	if (filter)
		return PyUnicode_FromString(filter);
	Py_RETURN_NONE;
}

static PyObject *
CatFile_object_name(CatFile *self, void *data)
{
	const char *objname;

	if (!self->f)
		return err_closed();

	objname = cmpack_cat_gkys(self->f, "object");
	if (objname) 
		return PyUnicode_FromString(objname);
	Py_RETURN_NONE;
}

static PyObject *
CatFile_right_ascension(CatFile *self, void *data)
{
	const char *str;
	double val;

	if (!self->f)
		return err_closed();

	str = cmpack_cat_gkys(self->f, "ra2000");
	if (str) {
		if (cmpack_strtora(str, &val)==0)
			return PyFloat_FromDouble(val);
	}
	Py_RETURN_NONE;
}

static PyObject *
CatFile_declination(CatFile *self, void *data)
{
	const char *str;
	double val;

	if (!self->f)
		return err_closed();

	str = cmpack_cat_gkys(self->f, "dec2000");
	if (str) {
		if (cmpack_strtodec(str, &val)==0)
			return PyFloat_FromDouble(val);
	}
	Py_RETURN_NONE;
}

static PyObject *
CatFile_object(CatFile *self, void *data)
{
	PyObject *res = PyTuple_New(3);
	PyTuple_SET_ITEM(res, 0, CatFile_object_name(self, data));
	PyTuple_SET_ITEM(res, 1, CatFile_right_ascension(self, data));
	PyTuple_SET_ITEM(res, 2, CatFile_declination(self, data));
	return res;
}

static PyObject *
CatFile_location_name(CatFile *self, void *data)
{
	const char *location;

	if (!self->f)
		return err_closed();

	location = cmpack_cat_gkys(self->f, "observatory");
	if (location) 
		return PyUnicode_FromString(location);
	Py_RETURN_NONE;
}

static PyObject *
CatFile_longitude(CatFile *self, void *data)
{
	const char *str;
	double val;

	if (!self->f)
		return err_closed();

	str = cmpack_cat_gkys(self->f, "longitude");
	if (str) {
		if (cmpack_strtolon(str, &val)==0)
			return PyFloat_FromDouble(val);
	}
	Py_RETURN_NONE;
}

static PyObject *
CatFile_latitude(CatFile *self, void *data)
{
	const char *str;
	double val;

	if (!self->f)
		return err_closed();

	str = cmpack_cat_gkys(self->f, "latitude");
	if (str) {
		if (cmpack_strtolat(str, &val)==0)
			return PyFloat_FromDouble(val);
	}
	Py_RETURN_NONE;
}

static PyObject *
CatFile_location(CatFile *self, void *data)
{
	PyObject *res = PyTuple_New(3);
	PyTuple_SET_ITEM(res, 0, CatFile_location_name(self, data));
	PyTuple_SET_ITEM(res, 1, CatFile_longitude(self, data));
	PyTuple_SET_ITEM(res, 2, CatFile_latitude(self, data));
	return res;
}

static PyObject *
CatFile_stars(CatFile *self, void *data)
{
	int i, count;
	PyObject *res;
	CmpackCatObject obj;

	if (!self->f)
		return err_closed();

	count = cmpack_cat_nstar(self->f);
	res = PyList_New(count);
	for (i=0; i<count; i++) {
		cmpack_cat_get_star(self->f, i, CMPACK_OM_ID, &obj);
		PyList_SET_ITEM(res, i, PyLong_FromLong(obj.id));
	}
	return res;
}

static PyObject *
CatFile_get_star(CatFile *self, PyObject *args)
{
	int index, id;
	CmpackCatObject obj;

	if (!PyArg_ParseTuple(args, "i", &id))
		return NULL;

	if (!self->f)
		return err_closed();

	index = cmpack_cat_find_star(self->f, id);
	if (index<0)
		Py_RETURN_NONE;

	cmpack_cat_get_star(self->f, index, CMPACK_OM_MAGNITUDE | CMPACK_OM_CENTER, &obj);
	if (obj.refmag_valid)
		return Py_BuildValue("ddd", obj.center_x, obj.center_y, obj.refmagnitude);
	else
		return Py_BuildValue("ddO", obj.center_x, obj.center_y, Py_None);
}

static PyObject *
CatFile_set_star(CatFile *self, PyObject *args)
{
	int res, index, id;
	double x, y;
	PyObject *mag = Py_None;
	CmpackCatObject obj;

	if (!PyArg_ParseTuple(args, "idd|O", &id, &x, &y, &mag))
		return NULL;

	if (!self->f)
		return err_closed();

	obj.id = id;
	obj.center_x = x;
	obj.center_y = y;
	obj.refmag_valid = PyFloat_Check(mag);
	if (obj.refmag_valid)
		obj.refmagnitude = PyFloat_AsDouble(mag);		

	index = cmpack_cat_find_star(self->f, id);
	if (index<0) 
		res = cmpack_cat_add_star(self->f, CMPACK_OM_MAGNITUDE | CMPACK_OM_CENTER, &obj);
	else
		res = cmpack_cat_set_star(self->f, index, CMPACK_OM_MAGNITUDE | CMPACK_OM_CENTER, &obj);
	if (res!=0)
		return cmpack_error(res);
	Py_RETURN_TRUE;
}

static PyObject *
CatFile_selection(CatFile *self, void *data)
{
	CmpackSelectionType type, t;
	int i, j, id, count, nstars;
	PyObject *res, *list;

	if (!self->f)
		return err_closed();

	res = PyDict_New();
	for (type=CMPACK_SELECT_VAR; type<CMPACK_SELECT_COUNT; type++) {
		nstars = cmpack_cat_get_selection_count(self->f);
		for (i=count=0; i<nstars; i++) {
			cmpack_cat_get_selection(self->f, i, NULL, &t);
			if (type == t) 
				count++;
		}
		list = PyList_New(count);
		for (i=j=0; i<nstars; i++) {
			cmpack_cat_get_selection(self->f, i, &id, &t);
			if (type == t)
				PyList_SET_ITEM(list, j++, PyLong_FromLong(id));
		}
		PyDict_SetItem(res, PyLong_FromLong(type), list);
	}
	return res;
}

static PyObject *
CatFile_get_selection(CatFile *self, PyObject *args)
{
	int i, j, id, count, nstars;
	CmpackSelectionType type, t;
	PyObject *res;

	if (!PyArg_ParseTuple(args, "i", &type))
		return NULL;

	if (!self->f)
		return err_closed();

	nstars = cmpack_cat_get_selection_count(self->f);
	for (i=count=0; i<nstars; i++) {
		cmpack_cat_get_selection(self->f, i, NULL, &t);
		if (type == t) 
			count++;
	}
	res = PyList_New(count);
	for (i=j=0; i<nstars; i++) {
		cmpack_cat_get_selection(self->f, i, &id, &t);
		if (type == t)
			PyList_SET_ITEM(res, j++, PyLong_FromLong(id));
	}
	return res;
}

static PyObject *
CatFile_clear_selection(CatFile *self, PyObject *args)
{
	if (!self->f)
		return err_closed();

	cmpack_cat_clear_selection(self->f);
	Py_RETURN_TRUE;
}

static PyObject *
CatFile_select(CatFile *self, PyObject *args)
{
	int id;
	CmpackSelectionType type;

	if (!PyArg_ParseTuple(args, "ii", &id, &type))
		return NULL;

	if (!self->f)
		return err_closed();

	cmpack_cat_update_selection(self->f, id, type);
	Py_RETURN_TRUE;
}

static PyObject *
CatFile_unselect(CatFile *self, PyObject *args)
{
	int id;

	if (!PyArg_ParseTuple(args, "i", &id))
		return NULL;

	if (!self->f)
		return err_closed();

	cmpack_cat_update_selection(self->f, id, CMPACK_SELECT_NONE);
	Py_RETURN_TRUE;
}

static PyObject *
CatFile_tags(CatFile *self, void *data)
{
	int i, star_id, count;
	const char *tag;
	PyObject *res;

	if (!self->f)
		return err_closed();

	res = PyDict_New();
	count = cmpack_cat_get_tag_count(self->f);
	for (i=0; i<count; i++) {
		cmpack_cat_get_tag(self->f, i, &star_id, &tag);
		PyDict_SetItem(res, PyLong_FromLong(star_id), PyUnicode_FromString(tag));
	}
	return res;
}

static PyObject *
CatFile_clear_tags(CatFile *self, PyObject *args)
{
	if (!self->f)
		return err_closed();

	cmpack_cat_clear_tags(self->f);
	Py_RETURN_TRUE;
}

static PyObject *
CatFile_update_tag(CatFile *self, PyObject *args)
{
	const char *tag;
	int id;

	if (!PyArg_ParseTuple(args, "iz", &id, &tag))
		return NULL;

	if (!self->f)
		return err_closed();

	cmpack_cat_update_tag(self->f, id, tag);
	Py_RETURN_TRUE;
}

static PyObject *
CatFile_remove_tag(CatFile *self, PyObject *args)
{
	int id;

	if (!PyArg_ParseTuple(args, "i", &id))
		return NULL;

	if (!self->f)
		return err_closed();

	cmpack_cat_remove_tag(self->f, id);
	Py_RETURN_TRUE;
}

static PyGetSetDef CatFile_getset[] = {
	{"image_size", (getter)CatFile_image_size, NULL,
	"Image width and height in pixels"},
	{"date_time", (getter)CatFile_date_time, NULL,
	"Date and time of observation"},
	{"julian_date", (getter)CatFile_julian_date, NULL,
	"Julian date of observation"},
	{"exposure", (getter)CatFile_exposure, NULL,
	"Exposure duration in seconds"},
	{"filter", (getter)CatFile_filter, NULL,
	"Color filter designation"},
	{"observer", (getter)CatFile_observer, NULL,
	"Observer's name"},
	{"object", (getter)CatFile_object, NULL, 
	"Object designation and coordinates"},
	{"location", (getter)CatFile_location, NULL,
	"Observatory designation and coordinates"},
	{"stars", (getter)CatFile_stars, NULL,
	"List of star identifiers"},
	{"selection", (getter)CatFile_selection, NULL,
	"Dictionary that maps selection type to a list of objects of that type"},
	{"tags", (getter)CatFile_tags, NULL,
	"Dictionary that maps object identifiers to tags assigned to them"},
	{NULL}
};

static PyMethodDef CatFile_methods[] = {
	{"test", (PyCFunction)CatFile_test, METH_VARARGS | METH_STATIC,
	"Check if the file is a Cat frame"},
	{"open", (PyCFunction)CatFile_open, METH_VARARGS | METH_STATIC,
	"Open a CCD frame stored in a file"},

	{"close", (PyCFunction)CatFile_close, METH_NOARGS,
	"Close file"},
	{"copy", (PyCFunction)CatFile_copy, METH_NOARGS,
	"Make deep copy of the frame"},
	{"save", (PyCFunction)CatFile_save, METH_KEYWORDS,
	"Save frame to a file"},

    {"__enter__", (PyCFunction)CatFile_self, METH_NOARGS,  
	"Checks if the file is open."},
    {"__exit__",  (PyCFunction)CatFile_exit, METH_VARARGS, 
	"Closes the file."}, 

	{"get_star", (PyCFunction)CatFile_get_star, METH_VARARGS,
	"Get star attributes"},
	{"set_star", (PyCFunction)CatFile_set_star, METH_VARARGS,
	"Set star attributes"},
	{"get_selection", (PyCFunction)CatFile_get_selection, METH_VARARGS,
	"Get list of stars of specified type"},
	{"clear_selection", (PyCFunction)CatFile_clear_selection, METH_NOARGS,
	"Update list of stars of specified type"},
	{"select", (PyCFunction)CatFile_select, METH_VARARGS,
	"Add a star to a list of specified type"},
	{"unselect", (PyCFunction)CatFile_unselect, METH_VARARGS,
	"Remove a star from all selection lists"},
	{"clear_tags", (PyCFunction)CatFile_clear_tags, METH_NOARGS,
	"Remove all object tags"},
	{"update_tag", (PyCFunction)CatFile_update_tag, METH_VARARGS,
	"Add or update tag for the specified object"},
	{"remove_tag", (PyCFunction)CatFile_remove_tag, METH_VARARGS,
	"Remove tag from the specified object"},

    {NULL}  /* Sentinel */
};

PyTypeObject CatFileType = {
    PyObject_HEAD_INIT(NULL)
    .tp_name = "cmpack.CatFile",          /*tp_name*/
    .tp_basicsize = sizeof(CatFile),           /*tp_basicsize*/
    .tp_dealloc = (destructor)CatFile_dealloc,           /*tp_dealloc*/
    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,        /*tp_flags*/
    .tp_doc = "Catalog file",        /* tp_doc */
    .tp_methods = CatFile_methods,           /* tp_methods */
    .tp_getset = CatFile_getset,            /* tp_getset */
    .tp_init = (initproc)CatFile_init,      /* tp_init */
    .tp_new = CatFile_new                 /* tp_new */
};

PyObject *
CatFile_FromHandle(CmpackCatFile *f)
{
	CatFile *pDoc = PyObject_New(CatFile, &CatFileType);
	if (pDoc)
		pDoc->f = f;
	return (PyObject*)pDoc;
}


PyTypeObject *CatFile_Type(void)
{
	return &CatFileType;
}
