#include "visualizer/renderer.h"

#include <assert.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "visualizer/grid.h"

static void append_repeat(struct PsiStringBuilder* sb, const char* glyph, size_t n)
{
    for (size_t i = 0; i < n; i++)
        psi_string_builder_append(sb, glyph);
}

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 void single_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;
        default: snprintf(out, cap, "[?]"); break;
    }
}

static void controlled_label(struct PsiGateOp op, char* out, size_t cap)
{
    switch (op.kind)
    {
        case PSI_GATE_CRX: snprintf(out, cap, "[CRx(%.2f)]", op.params[0]); break;
        case PSI_GATE_CRY: snprintf(out, cap, "[CRy(%.2f)]", op.params[0]); break;
        case PSI_GATE_CRZ: snprintf(out, cap, "[CRz(%.2f)]", op.params[0]); break;
        case PSI_GATE_CP: snprintf(out, cap, "[CP(%.2f)]", op.params[0]); break;
        default: snprintf(out, cap, "[?]"); break;
    }
}

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_horizontal(const struct PsiQuantumCircuit* circuit)
{
    size_t nq = circuit->num_qubits;
    size_t nc = circuit->num_classical;

    struct PsiStringBuilder* q = nq > 0 ? malloc(nq * sizeof(struct PsiStringBuilder)) : NULL;
    struct PsiStringBuilder* c = nc > 0 ? malloc(nc * sizeof(struct PsiStringBuilder)) : NULL;
    assert(q != NULL || nq == 0);
    assert(c != NULL || nc == 0);

    size_t max_label = 3;
    for (size_t i = 0; i < nq; i++)
    {
        char lbl[16];
        snprintf(lbl, sizeof lbl, "q%zu: ", i);
        if (strlen(lbl) > max_label)
            max_label = strlen(lbl);
    }
    for (size_t i = 0; i < nc; i++)
    {
        char lbl[16];
        snprintf(lbl, sizeof lbl, "c%zu: ", i);
        if (strlen(lbl) > max_label)
            max_label = strlen(lbl);
    }

    for (size_t i = 0; i < nq; i++)
    {
        q[i] = psi_new_string_builder();
        char lbl[16];
        snprintf(lbl, sizeof lbl, "q%zu: ", i);
        append_repeat(&q[i], " ", max_label - strlen(lbl));
        psi_string_builder_append(&q[i], lbl);
    }
    for (size_t i = 0; i < nc; i++)
    {
        c[i] = psi_new_string_builder();
        char lbl[16];
        snprintf(lbl, sizeof lbl, "c%zu: ", i);
        append_repeat(&c[i], " ", max_label - strlen(lbl));
        psi_string_builder_append(&c[i], lbl);
    }

    struct PsiStringBuilder gap = psi_new_string_builder();
    append_repeat(&gap, " ", max_label);

    if (circuit->operation_count == 0)
    {
        for (size_t i = 0; i < nq; i++)
            psi_string_builder_append(&q[i], "───");
        for (size_t i = 0; i < nc; i++)
            psi_string_builder_append(&c[i], "═══");
        psi_string_builder_append(&gap, "   ");
    }

    for (size_t oi = 0; oi < circuit->operation_count; oi++)
    {
        struct PsiGateOp op = circuit->operations[oi];

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

        if (is_single_target(op.kind))
        {
            char label[64];
            single_label(op, label, sizeof label);
            size_t w = psi_utf8_count(label);
            size_t t = op.qubits[0];

            for (size_t i = 0; i < nq; i++)
            {
                if (i == t)
                {
                    psi_string_builder_append(&q[i], "─");
                    psi_string_builder_append(&q[i], label);
                    psi_string_builder_append(&q[i], "─");
                }
                else
                    append_repeat(&q[i], "─", w + 2);
            }
            for (size_t i = 0; i < nc; i++)
                append_repeat(&c[i], "═", w + 2);
            append_repeat(&gap, " ", w + 2);
        }
        else if (is_param_controlled(op.kind))
        {
            char label[64];
            controlled_label(op, label, sizeof label);
            size_t w = psi_utf8_count(label);
            size_t control = op.qubits[0];
            size_t target = op.qubits[1];

            for (size_t i = 0; i < nq; i++)
            {
                psi_string_builder_append(&q[i], "─");
                if (i == control)
                {
                    psi_string_builder_append(&q[i], "●");
                    append_repeat(&q[i], "─", w - 1);
                }
                else if (i == target)
                    psi_string_builder_append(&q[i], label);
                else if (i > min_q && i < max_q)
                {
                    psi_string_builder_append(&q[i], "│");
                    append_repeat(&q[i], "─", w - 1);
                }
                else
                    append_repeat(&q[i], "─", w);
                psi_string_builder_append(&q[i], "─");
            }
            for (size_t i = 0; i < nc; i++)
                append_repeat(&c[i], "═", w + 2);
            append_repeat(&gap, " ", w + 2);
        }
        else if (op.kind == PSI_GATE_CNOT || op.kind == PSI_GATE_CZ || op.kind == PSI_GATE_SWAP ||
                 op.kind == PSI_GATE_CCNOT || op.kind == PSI_GATE_CSWAP)
        {
            for (size_t i = 0; i < nq; i++)
            {
                const char* seg = "─────";

                switch (op.kind)
                {
                    case PSI_GATE_CNOT:
                        if (i == op.qubits[0])
                            seg = "──●──";
                        else if (i == op.qubits[1])
                            seg = "──⊕──";
                        else if (i > min_q && i < max_q)
                            seg = "──│──";
                        break;
                    case PSI_GATE_CZ:
                        if (i == op.qubits[0] || i == op.qubits[1])
                            seg = "──●──";
                        else if (i > min_q && i < max_q)
                            seg = "──│──";
                        break;
                    case PSI_GATE_SWAP:
                        if (i == op.qubits[0] || i == op.qubits[1])
                            seg = "──╳──";
                        else if (i > min_q && i < max_q)
                            seg = "──│──";
                        break;
                    case PSI_GATE_CCNOT:
                        if (i == op.qubits[0] || i == op.qubits[1])
                            seg = "──●──";
                        else if (i == op.qubits[2])
                            seg = "──⊕──";
                        else if (i > min_q && i < max_q)
                            seg = "──│──";
                        break;
                    case PSI_GATE_CSWAP:
                        if (i == op.qubits[0])
                            seg = "──●──";
                        else if (i == op.qubits[1] || i == op.qubits[2])
                            seg = "──╳──";
                        else if (i > min_q && i < max_q)
                            seg = "──│──";
                        break;
                    default: break;
                }

                psi_string_builder_append(&q[i], seg);
            }
            for (size_t i = 0; i < nc; i++)
                psi_string_builder_append(&c[i], "═════");
            psi_string_builder_append(&gap, "     ");
        }
        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++)
            {
                if (i == mq)
                    psi_string_builder_append(&q[i], "─[M]─");
                else if (i > mq)
                    psi_string_builder_append(&q[i], "──║──");
                else
                    psi_string_builder_append(&q[i], "─────");
            }
            for (size_t i = 0; i < nc; i++)
            {
                if (i == mc)
                    psi_string_builder_append(&c[i], "══╩══");
                else if (i < mc)
                    psi_string_builder_append(&c[i], "══║══");
                else
                    psi_string_builder_append(&c[i], "═════");
            }
            psi_string_builder_append(&gap, "  ║  ");
        }
        else
        {
            char label[64];
            snprintf(label, sizeof label, "[%s]", op.custom->name);
            size_t w = psi_utf8_count(label);

            for (size_t i = 0; i < nq; i++)
            {
                bool is_target = false;
                for (size_t k = 0; k < tc; k++)
                    if (targets[k] == i)
                        is_target = true;

                psi_string_builder_append(&q[i], "─");
                if (is_target && i == targets[0])
                    psi_string_builder_append(&q[i], label);
                else if (is_target)
                    append_repeat(&q[i], "─", w);
                else if (i > min_q && i < max_q)
                {
                    psi_string_builder_append(&q[i], "│");
                    append_repeat(&q[i], "─", w - 1);
                }
                else
                    append_repeat(&q[i], "─", w);
                psi_string_builder_append(&q[i], "─");
            }
            for (size_t i = 0; i < nc; i++)
                append_repeat(&c[i], "═", w + 2);
            append_repeat(&gap, " ", w + 2);
        }
    }

    struct PsiStringBuilder out = psi_new_string_builder();

    for (size_t i = 0; i < nq; i++)
    {
        psi_string_builder_append(&out, psi_string_builder_finish(&q[i]));
        psi_string_builder_append(&out, "░\n");
        free(q[i].data);
    }

    if (nc > 0)
    {
        psi_string_builder_append(&out, psi_string_builder_finish(&gap));
        psi_string_builder_append(&out, "░\n");
        for (size_t i = 0; i < nc; i++)
        {
            psi_string_builder_append(&out, psi_string_builder_finish(&c[i]));
            psi_string_builder_append(&out, "░\n");
            free(c[i].data);
        }
    }

    free(gap.data);
    free(q);
    free(c);

    return psi_string_builder_finish(&out);
}