#include "postgres.h"

#include <sys/file.h>
#include <sys/stat.h>
#include <unistd.h>

#include "access/genam.h"
#include "access/heapam.h"
#include "access/htup.h"
#include "catalog/catalog.h"
#include "catalog/indexing.h"
#include "catalog/pg_class.h"
#include "fmgr.h"
#include "miscadmin.h"
#include "storage/relfilenode.h"
#include "utils/fmgroids.h"
#include "utils/syscache.h"


PG_MODULE_MAGIC;

Datum pg_relation_size_nolock_oid(PG_FUNCTION_ARGS);
PG_FUNCTION_INFO_V1(pg_relation_size_nolock_oid);

Datum pg_total_relation_size_nolock_oid(PG_FUNCTION_ARGS);
PG_FUNCTION_INFO_V1(pg_total_relation_size_nolock_oid);

static RelFileNode *
rfn_from_oid(Oid relid)
{
	Form_pg_class	relform;
	HeapTuple		tuple;
	Oid				filenode;
	RelFileNode	   *rfn;

	tuple = SearchSysCache(RELOID, ObjectIdGetDatum(relid), 0, 0, 0);
	if (!HeapTupleIsValid(tuple))
		elog(ERROR, "cache lookup failed for relation %u", relid);
	relform = (Form_pg_class) GETSTRUCT(tuple);

	switch (relform->relkind)
	{
		case RELKIND_RELATION:
		case RELKIND_INDEX:
		case RELKIND_SEQUENCE:
		case RELKIND_TOASTVALUE:
			/* okay, these have storage */
			filenode = relform->relfilenode;
			break;
		default:
			/* shouldn't have been called */
			elog(ERROR, "cannot calculate size of relation of kind %c",
				 relform->relkind);
			filenode = InvalidOid;	/* keep compiler quiet */
	}

	if (!OidIsValid(filenode))
		elog(ERROR, "relation %u does not have a valid relfilenode",
			 relid);

	rfn = palloc0(sizeof(RelFileNode));

	/* file proper: easy -- from pg_class */
	rfn->relNode = filenode;

	/*
	 * tablespace: if set on pg_class, use that.  Otherwise use the
	 * database default (which is always set, see InitPostgres)
	 */
	if (OidIsValid(relform->reltablespace))
		rfn->spcNode = relform->reltablespace;
	else
		rfn->spcNode = MyDatabaseTableSpace;

	/* database: is shared, then Invalid; otherwise, current database */
	if (relform->relisshared)
		rfn->dbNode = InvalidOid;
	else
		rfn->dbNode = MyDatabaseId;

	ReleaseSysCache(tuple);

	return rfn;
}

static int64
calculate_relation_size_internal(RelFileNode *rfn)
{
	int64		totalsize = 0;
	char	   *relationpath;
	char		pathname[MAXPGPATH];
	unsigned int segcount = 0;

	relationpath = relpath(*rfn);

	for (segcount = 0;; segcount++)
	{
		struct stat fst;

		if (segcount == 0)
			snprintf(pathname, MAXPGPATH, "%s",
					 relationpath);
		else
			snprintf(pathname, MAXPGPATH, "%s.%u",
					 relationpath, segcount);

		if (stat(pathname, &fst) < 0)
		{
			if (errno == ENOENT)
				break;
			else
				ereport(ERROR,
						(errcode_for_file_access(),
						 errmsg("could not stat file \"%s\": %m", pathname)));
		}
		totalsize += fst.st_size;
	}

	return totalsize;
}

/* Given an OID, returns the size for the bare relation identified by it */
static int64
calculate_relation_size(Oid relid)
{
	RelFileNode	*rfn;
	int64		size;

	rfn = rfn_from_oid(relid);
	size = calculate_relation_size_internal(rfn);
	pfree(rfn);

	return size;
}

Datum
pg_relation_size_nolock_oid(PG_FUNCTION_ARGS)
{
	Oid			relid = PG_GETARG_OID(0);
	int64		size;

	size = calculate_relation_size(relid);

	PG_RETURN_INT64(size);
}

/*
 * Return a table's index list, as an OID list
 * 		This is roughly equivalent to RelationGetIndexList, except we only want
 * 		an OID as output, and we avoid locks on both the table and indexes.
 */
static List *
OidGetIndexList(Oid relid)
{
	Relation	indrel;
	SysScanDesc	indscan;
	ScanKeyData	skey;
	HeapTuple	htup;
	List	   *result = NIL;

	ScanKeyInit(&skey,
				Anum_pg_index_indrelid,
				BTEqualStrategyNumber, F_OIDEQ,
				ObjectIdGetDatum(relid));

	indrel = heap_open(IndexRelationId, AccessShareLock);
	indscan = systable_beginscan(indrel, IndexIndrelidIndexId, true,
								 SnapshotNow, 1, &skey);

	while (HeapTupleIsValid(htup = systable_getnext(indscan)))
	{
		Form_pg_index	indform;

		indform = (Form_pg_index) GETSTRUCT(htup);
		result = lappend_oid(result, indform->indexrelid);
	}

	systable_endscan(indscan);
	heap_close(indrel, AccessShareLock);

	return result;
}

/*
 * returns sum of sizes of the relation.
 *
 * For a plain relation, adds sizes of indexes, and recurses to add
 * size of toast table and index.
 *
 * For a TOAST table, adds size of its index
 */
static int64
total_relation_size_internal(Oid relid)
{
	int64		size;
	HeapTuple	tuple;
	Form_pg_class	classtup;

	tuple = SearchSysCache(RELOID, ObjectIdGetDatum(relid), 0, 0, 0);
	if (!HeapTupleIsValid(tuple))
		elog(ERROR, "cache lookup failed for relation %u", relid);
	classtup = (Form_pg_class) GETSTRUCT(tuple);

	size = calculate_relation_size(relid);

	/*
	 * for a heap we also need to consider its indexes, and also the toast
	 * table and respective index.
	 */
	if (classtup->relkind == RELKIND_RELATION)
	{
		List	*idxs;
		ListCell *cell;

		idxs = OidGetIndexList(relid);
		foreach (cell, idxs)
		{
			Oid		idxid = lfirst_oid(cell);

			size += calculate_relation_size(idxid);
		}
		list_free(idxs);

		/* recurse to add data for TOAST items */
		if (classtup->reltoastrelid)
			size += total_relation_size_internal(classtup->reltoastrelid);
	}

	/* for a TOAST table, add the toast index */
	if (classtup->relkind == RELKIND_TOASTVALUE)
	{
		Assert(classtup->reltoastidxid);
		
		size += calculate_relation_size(classtup->reltoastidxid);
	}

	ReleaseSysCache(tuple);

	return size;
}

Datum
pg_total_relation_size_nolock_oid(PG_FUNCTION_ARGS)
{
	Oid		relid = PG_GETARG_OID(0);
	int64	size;

	size = total_relation_size_internal(relid);

	PG_RETURN_INT64(size);
}