#include "visualizer/renderer.h"

#include <stdbool.h>
#include <stdio.h>

#include "visualizer/grid.h"

static void gate_label(struct PsiGateOp op, char *out, size_t cap)
{
	switch (op.kind)
	{
		case PSI_GATE_H: snprintf(out, cap, "[H]"); break;
		case PSI_GATE_X: snprintf(out, cap, "[X]"); break;
		case PSI_GATE_Y: snprintf(out, cap, "[Y]"); break;
		case PSI_GATE_Z: snprintf(out, cap, "[Z]"); break;
		case PSI_GATE_S: snprintf(out, cap, "[S]"); break;
		case PSI_GATE_T: snprintf(out, cap, "[T]"); break;
		case PSI_GATE_SDG: snprintf(out, cap, "[S†]"); break;
		case PSI_GATE_TDG: snprintf(out, cap, "[T†]"); break;
		case PSI_GATE_SX: snprintf(out, cap, "[√X]"); break;
		case PSI_GATE_SXDG: snprintf(out, cap, "[√X†]"); break;
		case PSI_GATE_RX: snprintf(out, cap, "[Rx(%.2f)]", op.params[0]); break;
		case PSI_GATE_RY: snprintf(out, cap, "[Ry(%.2f)]", op.params[0]); break;
		case PSI_GATE_RZ: snprintf(out, cap, "[Rz(%.2f)]", op.params[0]); break;
		case PSI_GATE_P: snprintf(out, cap, "[P(%.2f)]", op.params[0]); break;
		case PSI_GATE_U1: snprintf(out, cap, "[U1(%.2f)]", op.params[0]); break;
		case PSI_GATE_U2: snprintf(out, cap, "[U2]"); break;
		case PSI_GATE_U3: snprintf(out, cap, "[U3]"); break;
		case PSI_GATE_CRX: snprintf(out, cap, "[CRx]"); break;
		case PSI_GATE_CRY: snprintf(out, cap, "[CRy]"); break;
		case PSI_GATE_CRZ: snprintf(out, cap, "[CRz]"); break;
		case PSI_GATE_CP: snprintf(out, cap, "[CP]"); break;
		case PSI_GATE_CNOT: snprintf(out, cap, "●"); break;
		case PSI_GATE_CZ: snprintf(out, cap, "●"); break;
		case PSI_GATE_SWAP: snprintf(out, cap, "╳"); break;
		case PSI_GATE_CCNOT: snprintf(out, cap, "●"); break;
		case PSI_GATE_CSWAP: snprintf(out, cap, "●"); break;
		case PSI_GATE_MEASURE: snprintf(out, cap, "[M]"); break;
		case PSI_GATE_CUSTOM: snprintf(out, cap, "[%s]", op.custom->name); break;
	}
}

static size_t calculate_col_width(const struct PsiQuantumCircuit *circuit)
{
	size_t max_label = 3;

	for (size_t i = 0; i < circuit->operation_count; i++)
	{
		char label[64];
		gate_label(circuit->operations[i], label, sizeof label);
		size_t n = psi_utf8_count(label);
		if (n > max_label)
			max_label = n;
	}

	size_t width = max_label + 2;
	if (width % 2 == 0)
		return width + 1;

	return width;
}

static bool is_single_target(enum PsiGateKind kind)
{
	switch (kind)
	{
		case PSI_GATE_H:
		case PSI_GATE_X:
		case PSI_GATE_Y:
		case PSI_GATE_Z:
		case PSI_GATE_S:
		case PSI_GATE_T:
		case PSI_GATE_SDG:
		case PSI_GATE_TDG:
		case PSI_GATE_SX:
		case PSI_GATE_SXDG:
		case PSI_GATE_RX:
		case PSI_GATE_RY:
		case PSI_GATE_RZ:
		case PSI_GATE_P:
		case PSI_GATE_U1:
		case PSI_GATE_U2:
		case PSI_GATE_U3:
			return true;
		default:
			return false;
	}
}

static bool is_param_controlled(enum PsiGateKind kind)
{
	switch (kind)
	{
		case PSI_GATE_CRX:
		case PSI_GATE_CRY:
		case PSI_GATE_CRZ:
		case PSI_GATE_CP:
			return true;
		default:
			return false;
	}
}

char *psi_render_circuit_vertical(const struct PsiQuantumCircuit *circuit)
{
	struct PsiStringBuilder sb = psi_new_string_builder();

	size_t nq = circuit->num_qubits;
	size_t nc = circuit->num_classical;

	if (nq == 0)
		return psi_string_builder_finish(&sb);

	size_t col_width = calculate_col_width(circuit);
	size_t gap_width = 3;
	size_t stride = col_width + 1;
	size_t half = col_width / 2;

	size_t q_total = nq * col_width + (nq - 1);
	size_t c_total = nc > 0 ? nc * col_width + (nc - 1) : 0;
	size_t total_width = q_total + gap_width + c_total;

	struct PsiGlyphRow header = psi_new_glyph_row(total_width);
	for (size_t i = 0; i < nq; i++)
	{
		char label[16];
		snprintf(label, sizeof label, "q%zu", i);
		size_t col_start = i * stride;
		psi_glyph_row_place(&header, col_start + (col_width - psi_utf8_count(label)) / 2, label);
	}
	for (size_t i = 0; i < nc; i++)
	{
		char label[16];
		snprintf(label, sizeof label, "c%zu", i);
		size_t col_start = q_total + gap_width + i * stride;
		psi_glyph_row_place(&header, col_start + (col_width - psi_utf8_count(label)) / 2, label);
	}
	psi_glyph_row_render(header, &sb);
	psi_free_glyph_row(&header);

	for (size_t op_index = 0; op_index <= circuit->operation_count; op_index++)
	{
		struct PsiGlyphRow wires = psi_new_glyph_row(total_width);
		for (size_t i = 0; i < nq; i++)
			psi_glyph_row_set(&wires, i * stride + half, "│");
		for (size_t i = 0; i < nc; i++)
			psi_glyph_row_set(&wires, q_total + gap_width + i * stride + half, "║");
		psi_glyph_row_render(wires, &sb);
		psi_free_glyph_row(&wires);

		if (op_index == circuit->operation_count)
			break;

		struct PsiGateOp op = circuit->operations[op_index];

		size_t target_count;
		const size_t *targets = psi_gate_op_quantum_targets(&op, &target_count);
		size_t min_q = targets[0];
		size_t max_q = targets[0];
		for (size_t i = 1; i < target_count; i++)
		{
			if (targets[i] < min_q)
				min_q = targets[i];
			if (targets[i] > max_q)
				max_q = targets[i];
		}

		char label[64];
		gate_label(op, label, sizeof label);
		size_t label_len = psi_utf8_count(label);

		struct PsiGlyphRow row = psi_new_glyph_row(total_width);

		if (is_single_target(op.kind))
		{
			size_t target = op.qubits[0];
			for (size_t i = 0; i < nq; i++)
			{
				size_t col_start = i * stride;
				if (i == target)
					psi_glyph_row_place(&row, col_start + (col_width - label_len) / 2, label);
				else
					psi_glyph_row_set(&row, col_start + half, "│");
			}
			for (size_t i = 0; i < nc; i++)
				psi_glyph_row_set(&row, q_total + gap_width + i * stride + half, "║");
		}
		else if (op.kind == PSI_GATE_MEASURE)
		{
			size_t mq = op.qubits[0];
			size_t mc = op.classical;
			for (size_t i = 0; i < nq; i++)
			{
				size_t col_start = i * stride;
				if (i < mq)
					psi_glyph_row_set(&row, col_start + half, "│");
				else if (i == mq)
					psi_glyph_row_place(&row, col_start + (col_width - label_len) / 2, label);
			}

			size_t mq_center = mq * stride + half;
			size_t mc_start = q_total + gap_width;
			size_t mc_center = mc_start + mc * stride + half;

			psi_glyph_row_fill_space(&row, mq_center + 2, mc_center + 1, "═");
			psi_glyph_row_set(&row, mc_center, "╣");

			for (size_t i = 0; i < nc; i++)
				if (i > mc)
					psi_glyph_row_set(&row, mc_start + i * stride + half, "║");
		}
		else if (op.kind == PSI_GATE_CCNOT || op.kind == PSI_GATE_CSWAP)
		{
			bool is_cswap = op.kind == PSI_GATE_CSWAP;
			const char *sym_c = "●";
			const char *sym_t = is_cswap ? "╳" : "⊕";
			size_t c1 = op.qubits[0];
			size_t c2 = op.qubits[1];
			size_t t = op.qubits[2];

			for (size_t i = 0; i < nq; i++)
			{
				size_t center = i * stride + half;
				if (i < min_q || i > max_q)
					psi_glyph_row_set(&row, center, "│");
				else if (i == c1)
					psi_glyph_row_set(&row, center, sym_c);
				else if (i == c2)
					psi_glyph_row_set(&row, center, is_cswap ? sym_t : sym_c);
				else if (i == t)
					psi_glyph_row_set(&row, center, sym_t);
			}

			psi_glyph_row_fill_space(&row, min_q * stride + half + 1, max_q * stride + half, "─");
			for (size_t i = 0; i < nc; i++)
				psi_glyph_row_set(&row, q_total + gap_width + i * stride + half, "║");
		}
		else if (op.kind == PSI_GATE_CUSTOM)
		{
			if (target_count == 1)
			{
				for (size_t i = 0; i < nq; i++)
				{
					size_t col_start = i * stride;
					if (i == targets[0])
						psi_glyph_row_place(&row, col_start + (col_width - label_len) / 2, label);
					else
						psi_glyph_row_set(&row, col_start + half, "│");
				}
			}
			else
			{
				for (size_t i = 0; i < nq; i++)
				{
					size_t col_start = i * stride;
					size_t center = col_start + half;
					bool is_target = false;
					for (size_t k = 0; k < target_count; k++)
						if (targets[k] == i)
							is_target = true;

					if (i < min_q || i > max_q)
						psi_glyph_row_set(&row, center, "│");
					else if (i == targets[0])
						psi_glyph_row_place(&row, col_start + (col_width - label_len) / 2, label);
					else if (is_target)
						psi_glyph_row_set(&row, center, "□");
				}

				psi_glyph_row_fill_space(&row, min_q * stride + half + 1, max_q * stride + half, "─");
			}

			for (size_t i = 0; i < nc; i++)
				psi_glyph_row_set(&row, q_total + gap_width + i * stride + half, "║");
		}
		else
		{
			const char *sym1 = "●";
			const char *sym2 = "⊕";
			if (op.kind == PSI_GATE_CZ)
				sym2 = "●";
			else if (op.kind == PSI_GATE_SWAP)
			{
				sym1 = "╳";
				sym2 = "╳";
			}
			else if (is_param_controlled(op.kind))
				sym2 = "□";

			size_t control = op.qubits[0];
			size_t target = op.qubits[1];

			for (size_t i = 0; i < nq; i++)
			{
				size_t col_start = i * stride;
				size_t center = col_start + half;
				if (i < min_q || i > max_q)
					psi_glyph_row_set(&row, center, "│");
				else if (i == control)
					psi_glyph_row_set(&row, center, sym1);
				else if (i == target)
				{
					if (is_param_controlled(op.kind))
						psi_glyph_row_place(&row, col_start + (col_width - label_len) / 2, label);
					else
						psi_glyph_row_set(&row, center, sym2);
				}
			}

			psi_glyph_row_fill_space(&row, min_q * stride + half + 1, max_q * stride + half, "─");
			for (size_t i = 0; i < nc; i++)
				psi_glyph_row_set(&row, q_total + gap_width + i * stride + half, "║");
		}

		psi_glyph_row_render(row, &sb);
		psi_free_glyph_row(&row);
	}

	struct PsiGlyphRow end_row = psi_new_glyph_row(total_width);
	for (size_t i = 0; i < total_width; i++)
		psi_glyph_row_set(&end_row, i, "░");
	psi_glyph_row_render(end_row, &sb);
	psi_free_glyph_row(&end_row);

	return psi_string_builder_finish(&sb);
}