import * as THREE from 'three';
import { CONFIG, STATE } from '../core/State.js';
import { createLabelSprite } from './LabelHelper.js';
import { vertexShader, fragmentShader } from '../engine/Shaders.js';
import { PALETTE_3D, getActivationColorPerceptual as getActivationColor } from '../engine/palette.js';

// ─── Helpers CPU ────────────────────────────────────────────────────────────
function lerp(a, b, t) { return a + (b - a) * t; }
// ─────────────────────────────────────────────────────────────────────────

export function initInstancedMesh(count, scene) {
    // Redirect to initMacroMode logic
    initMacroMode(scene);
    return null;
}

export function initMacroMode(scene) {
    if (STATE.macroMeshes) {
        if (STATE.macroMeshes.lines) scene.remove(STATE.macroMeshes.lines);
        if (STATE.macroMeshes.points) scene.remove(STATE.macroMeshes.points);
        STATE.macroMeshes = null;
    }
}

export function updateMacroCloud(tokens, scene) {
    if (!tokens || tokens.length === 0) return;

    if (STATE.macroMeshes) {
        if (STATE.macroMeshes.lines) {
            scene.remove(STATE.macroMeshes.lines);
            if (STATE.macroMeshes.lines.geometry) STATE.macroMeshes.lines.geometry.dispose();
            if (STATE.macroMeshes.lines.material) STATE.macroMeshes.lines.material.dispose();
        }
        if (STATE.macroMeshes.points) {
            scene.remove(STATE.macroMeshes.points);
            if (STATE.macroMeshes.points.geometry) STATE.macroMeshes.points.geometry.dispose();
            if (STATE.macroMeshes.points.material) STATE.macroMeshes.points.material.dispose();
        }
        STATE.macroMeshes = null;
    }

    // Cleanup Legacy InstancedMesh
    if (STATE.instancedMesh) {
        scene.remove(STATE.instancedMesh);
        if (STATE.instancedMesh.geometry) STATE.instancedMesh.geometry.dispose();
        if (STATE.instancedMesh.material) STATE.instancedMesh.material.dispose();
        STATE.instancedMesh = null;
    }

    // ─── Prepare Data Buffers ─────────────────────────────────────────────────
    let totalVertices = 0;
    tokens.forEach(t => {
        const len = t.embedding ? t.embedding.length : 0;
        if (len > 0) totalVertices += len;
    });

    const posArray     = new Float32Array(totalVertices * 3); // posiciones compartidas
    const lineColArray = new Float32Array(totalVertices * 3); // color dominante por thread (para la línea)
    const intensityArr = new Float32Array(totalVertices);     // valor de activación por vértice (para el shader de puntos)
    const vtxIdxArr    = new Float32Array(totalVertices);     // índice local por vértice (para highlight)
    const lineIndices  = [];

    const mode = STATE.currentZMode || 'log';
    const maxID = Math.max(...tokens.map(t => t.tokenID || t.id || 0), 1);

    const sortedTokens = [...tokens].sort((a, b) => (a.tokenID || a.id || 0) - (b.tokenID || b.id || 0));
    const rankMap = new Map();
    sortedTokens.forEach((t, i) => rankMap.set(t.id, i));

    let globalVertexIndex = 0;

    tokens.forEach((t, tIdx) => {
        const vec = t.vector || t.embedding;
        if (!vec || vec.length === 0) return;

        const len = vec.length;

        // Mantener color de identidad solo para labels
        if (!t.color) {
            const colorHex = CONFIG.identityColors[tIdx % CONFIG.identityColors.length];
            t.color = new THREE.Color(colorHex);
        }

        const rank  = rankMap.get(t.id) ?? 0;
        const idVal = t.tokenID || t.id || 0;

        let offsetX   = 0;
        let rotationZ = 0;
        const layoutMode = STATE.layoutMode || 'COMPARE';

        if (layoutMode === 'BLANKET') {
            if (mode === 'linear') {
                offsetX = rank * CONFIG.threadSpacing;
            } else if (mode === 'log') {
                const ratio = idVal / maxID;
                offsetX = ratio * 200.0 * CONFIG.zScale.log * CONFIG.threadSpacing;
            } else if (mode === 'radial') {
                rotationZ = rank * STATE.radialStep;
                offsetX   = rank * CONFIG.threadSpacing;
            }
        } else {
            offsetX = rank * (CONFIG.compareSpacing || 0.0);
            if (mode === 'radial') rotationZ = rank * STATE.radialStep;
        }

        // ─── Color dominante del thread: valor de mayor |activación| → color de la línea ───
        let maxAbsVal = 0, dominantVal = 0;
        for (let i = 0; i < len; i++) {
            if (Math.abs(vec[i]) > maxAbsVal) { maxAbsVal = Math.abs(vec[i]); dominantVal = vec[i]; }
        }
        const lineColor = getActivationColor(dominantVal);
        // Normalizador: garantiza que el valor máximo del token sea ±1.0 en el shader
        const safeMax = maxAbsVal > 0 ? maxAbsVal : 1.0;
        // ─────────────────────────────────────────────────────────────────────────

        for (let i = 0; i < len; i++) {
            const localY = vec[i] * CONFIG.amplitudeScale;
            const localZ = i * CONFIG.pointSpacing * (CONFIG.threadWidth || 1.0);

            let finalX = 0, finalY = 0, finalZ = 0;
            if (mode === 'radial') {
                const v = new THREE.Vector3(0, localY, localZ);
                v.applyEuler(new THREE.Euler(0, -Math.PI / 2, 0));
                v.applyAxisAngle(new THREE.Vector3(0, 0, 1), rotationZ);
                finalX = v.x; finalY = v.y; finalZ = v.z;
            } else {
                finalX = offsetX; finalY = localY; finalZ = localZ;
            }

            const ptr = globalVertexIndex * 3;
            posArray[ptr]     = finalX;
            posArray[ptr + 1] = finalY;
            posArray[ptr + 2] = finalZ;

            // Línea: color dominante (uniforme por thread)
            lineColArray[ptr]     = lineColor.r;
            lineColArray[ptr + 1] = lineColor.g;
            lineColArray[ptr + 2] = lineColor.b;

            // Puntos: valor de activación NORMALIZADO por vértice → el shader GLSL calcula el color
            // La normalización garantiza que los valores lleguen a ±1.0 independientemente de la escala del modelo
            intensityArr[globalVertexIndex] = vec[i] / safeMax;
            vtxIdxArr[globalVertexIndex]    = i;

            if (i < len - 1) lineIndices.push(globalVertexIndex, globalVertexIndex + 1);

            globalVertexIndex++;
        }

        // Labels (color de identidad se mantiene)
        if (STATE.scene) {
            const isHidden = STATE.hiddenThreads.has(tIdx);
            const tipY = vec[0] * CONFIG.amplitudeScale;
            let labelX = 0, labelY = 0, labelZ = 0;

            if (mode === 'radial') {
                const v = new THREE.Vector3(0, tipY, 0);
                v.applyEuler(new THREE.Euler(0, -Math.PI / 2, rotationZ));
                labelX = v.x; labelY = v.y; labelZ = v.z;
            } else {
                labelX = offsetX; labelY = tipY; labelZ = 0;
            }

            if (t.labelSprite) {
                t.labelSprite.position.set(labelX, labelY + 1.5, labelZ);
                t.labelSprite.visible = !isHidden;
            } else if (!isHidden) {
                const sprite = createLabelSprite(`${t.label || tIdx} [${idVal}]`, labelX, labelY + 1.5, labelZ, t.color);
                STATE.scene.add(sprite);
                t.labelSprite = sprite;
            }
        }
    });

    // ─── Geometría de Líneas: color dominante por thread ──────────────────────────
    const lineGeo = new THREE.BufferGeometry();
    lineGeo.setAttribute('position', new THREE.BufferAttribute(posArray,     3));
    lineGeo.setAttribute('color',    new THREE.BufferAttribute(lineColArray, 3));
    lineGeo.setIndex(lineIndices);

    const lineMat = new THREE.LineBasicMaterial({
        vertexColors: true,
        opacity: 0.55,
        transparent: true,
        linewidth: CONFIG.threadThickness || 1.0
    });

    // ─── Geometría de Puntos: ShaderMaterial con intensidad por vértice ───────────────
    const pointGeo = new THREE.BufferGeometry();
    pointGeo.setAttribute('position',    new THREE.BufferAttribute(posArray,     3));
    pointGeo.setAttribute('intensity',   new THREE.BufferAttribute(intensityArr, 1));
    pointGeo.setAttribute('vertexIndex', new THREE.BufferAttribute(vtxIdxArr,    1));

    const pointMat = new THREE.ShaderMaterial({
        uniforms: {
            highlightIndex: { value: -1.0 },
            pointSpacing:   { value: CONFIG.pointSpacing },
            // dotSize: tamaño visual del punto, escalado por amplitudeScale para compensar la escala del embedding
            dotSize:        { value: CONFIG.pointSpacing * CONFIG.amplitudeScale * 0.8 },
            opacity:        { value: 1.0 },
            color:          { value: new THREE.Color(PALETTE_3D.AMBIENT_LIGHT) }
        },
        vertexShader,
        fragmentShader,
        transparent: true,
        depthWrite:  false,
        blending:    THREE.NormalBlending
    });

    const lines  = new THREE.LineSegments(lineGeo,  lineMat);
    const points = new THREE.Points(pointGeo, pointMat);

    lines.frustumCulled  = false;
    points.frustumCulled = false;

    scene.add(lines);
    scene.add(points);

    STATE.macroMeshes = { lines, points };
}

export function updateInstancedCloud(tokens) {
    if (STATE.scene) { updateMacroCloud(tokens, STATE.scene); return; }
}