// src/components/OrganizationTreeDiagram.vue
<template>
  <div ref="diagramContainer" class="diagram-container"></div>
</template>

<script setup>
import { ref, onMounted, watch, onUnmounted } from 'vue';
import * as d3 from 'd3';
import { useAppStore } from '@/stores/app.js';

const appStore = useAppStore();

const props = defineProps({
  treeData: {
    type: Object, // Or Array if it's a list of roots
    required: true
  }
});

const diagramContainer = ref(null);

// Helper functions (unchanged)
const tailwindColors = {
  'blue-700': '#1d4ed8',
  'orange-500': '#f97316',
  'red-500': '#ef4444',
  'green-500': '#22c55e',
  'purple-600': '#9333ea',
  'blue-600': '#2563eb',
  'orange-400': '#fb923c',
  'red-400': '#f87171',
  'red-300': '#ef4444',
  'green-400': '#4ade80',
  'purple-500': '#a855f7',
  'gray-500': '#6b7280'
};

const tailwindTextColors = {
  'text-white': '#ffffff',
  'text-gray-800': '#1f2937'
};

const getNodeColorClass = (nodeId) => {
  if (nodeId === 'alpha_root') return 'bg-blue-700 text-white';
  if (nodeId === 'division_hq') return 'bg-orange-500 text-white';
  if (nodeId === 'operation_division') return 'bg-red-500 text-white';
  if (nodeId === 'support_division') return 'bg-green-500 text-white';
  if (nodeId === 'special_projects') return 'bg-purple-600 text-white';
  if (nodeId.startsWith('dept_A')) return 'bg-orange-400 text-white';
  if (nodeId.startsWith('section_B')) return 'bg-red-400 text-white';
  if (nodeId.startsWith('unit_B')) return 'bg-red-300 text-gray-800';
  if (nodeId.startsWith('sub_support_')) return 'bg-green-400 text-white';
  if (nodeId.startsWith('project_')) return 'bg-purple-500 text-white';
  if (nodeId === 'independent_unit_A' || nodeId === 'advisory_board') return 'bg-gray-500 text-white';
  return 'bg-blue-600 text-white';
};

const getNodeColor = (nodeId) => {
  const className = getNodeColorClass(nodeId);
  const bgColorClass = className.split(' ').find(cls => cls.startsWith('bg-'));
  return tailwindColors[bgColorClass.replace('bg-', '')] || '#666';
};

const getTextColor = (nodeId) => {
  const className = getNodeColorClass(nodeId);
  const textColorClass = className.split(' ').find(cls => cls.startsWith('text-'));
  return tailwindTextColors[textColorClass.replace('text-', '')] || '#fff';
};

// Define constants for node dimensions and spacing
const BASE_NODE_WIDTH = 180;
const BASE_NODE_HEIGHT = 80;
const HORIZONTAL_NODE_SPACING = 80;
const VERTICAL_DEPTH_SPACING = 150;

let svg, gChart;
let treemap;
let rootNode;
let i = 0; // Counter for assigning unique IDs

const setupD3Tree = (preserveState = false) => {
  console.log("setupD3Tree called.");
  if (!props.treeData || !diagramContainer.value) {
    console.warn("treeData or diagramContainer is null. Skipping setup.");
    return;
  }

  d3.select(diagramContainer.value).select('svg').remove();

  const containerWidth = diagramContainer.value.offsetWidth;
  const containerHeight = diagramContainer.value.offsetHeight;

  svg = d3.select(diagramContainer.value)
      .append("svg")
      .attr("width", containerWidth)
      .attr("height", containerHeight);

  gChart = svg.append("g");

  treemap = d3.tree().nodeSize([BASE_NODE_WIDTH + HORIZONTAL_NODE_SPACING, VERTICAL_DEPTH_SPACING]);

  const clonedTreeData = JSON.parse(JSON.stringify(props.treeData));

  // Handle array of roots by creating a dummy root
  const dataToHierarchy = Array.isArray(clonedTreeData)
      ? { id: 'dummy_root', name_th: '', name_en: '', children: clonedTreeData }
      : clonedTreeData;

  rootNode = d3.hierarchy(dataToHierarchy, d => d.children);

  console.log("D3 Hierarchy - Root Node children:", rootNode.children ? rootNode.children.map(c => c.id) : "No children (single root)");
  rootNode.descendants().forEach(d => {
    console.log(`D3 Hierarchy - Node: ${d.id}, Depth: ${d.depth}, isDashed: ${d.data.isDashed}, reportsTo: ${d.data.reportsTo}`);
  });

  treemap(rootNode); // Compute initial layout for all nodes

  rootNode.descendants().forEach(d => {
    d.id = d.data.id || ++i;
    d.data.rectWidth = BASE_NODE_WIDTH;
    d.data.rectHeight = BASE_NODE_HEIGHT;
    // Set initial position for animation
    d.x0 = d.x; // Use computed x for initial animation start
    d.y0 = d.y; // Use computed y for initial animation start
  });

  if (!preserveState) {
    console.log("Performing initial collapse.");
    rootNode.descendants().forEach(d => {
      // Adjust depth for array-based root, if dummy_root is depth 0
      const actualDepth = Array.isArray(props.treeData) ? d.depth - 1 : d.depth;
      // Collapse nodes at depth 2 or deeper, unless it's a dashed line node
      // Note: if you want to show root and direct children only, it's actualDepth >= 1
      if (d.children && actualDepth >= 1 && !d.data.isDashed) { // Collapse beyond depth 0 (root)
        d._children = d.children; // Store original children in _children
        d.children = null;       // Set children to null to hide them
      }
    });
  } else {
    console.log("Preserving state: Skipping initial collapse.");
  }

  update(rootNode); // Initial call to update function
};

// Helper to convert D3 transition to a Promise
function endAll(transition) {
  return new Promise(resolve => {
    if (transition.empty()) {
      resolve();
    } else {
      let n = transition.size();
      transition.on("end", () => {
        n--;
        if (n === 0) resolve();
      });
    }
  });
}

const update = (source) => {
  console.log("update called with source:", source.id);

  // Re-compute the layout with the current state of rootNode (collapsed/expanded)
  const treeData = treemap(rootNode);

  // Get all descendants (nodes) in the current hierarchy layout
  const allNodesInLayout = treeData.descendants();

  const isArrayRoot = Array.isArray(props.treeData);

  const visibleNodes = allNodesInLayout.filter(d => {

    if (d.id === 'dummy_root') return true;


    const isActualRoot = (!isArrayRoot && d.depth === 0) || (isArrayRoot && d.depth === 1 && d.parent.id === 'dummy_root');


    return isActualRoot || (d.parent && d.parent.children && d.parent.children.includes(d));
  });

  // Explicitly remove dummy_root from visibleNodes for rendering if it's an array root
  const nodesToRender = visibleNodes.filter(d => d.id !== 'dummy_root');
  console.log("Nodes to Render:", nodesToRender.map(d => d.id));


  let links = [];

  // Find the actual main root node (e.g., alpha_root) among all nodes in layout
  const mainRoot = allNodesInLayout.find(d => d.id === 'alpha_root');

  treemap(rootNode).links().forEach(link => {

    if (isArrayRoot && link.source.id === 'dummy_root') {
      if (link.target.data && !link.target.data.isDashed && nodesToRender.some(n => n.id === link.target.id)) {
        links.push(link);
      }
    } else {
      const sourceVisible = nodesToRender.some(n => n.id === link.source.id);
      const targetVisible = nodesToRender.some(n => n.id === link.target.id);
      if (sourceVisible && targetVisible) {
        links.push(link);
      }
    }
  });

  if (mainRoot) { // Only add dashed links if mainRoot exists
    allNodesInLayout.forEach(node => {
      if (node.data.isDashed && node.data.reportsTo === mainRoot.id) {
        const isNodeVisible = nodesToRender.some(n => n.id === node.id);
        const isSourceVisible = nodesToRender.some(n => n.id === mainRoot.id); // mainRoot should always be visible if it's the core.

        if (isNodeVisible && isSourceVisible) {

          links.push({
            source: mainRoot,
            target: node,
            data: {
              isDashed: true
            }
          });
        }
      }
    });
  }

  console.log("Final links (including custom):", links.map(d => {
    const sourceId = d.source ? d.source.id : 'N/A';
    const targetId = d.target ? d.target.id : 'N/A';
    const isDashed = (d.data && d.data.isDashed) || (d.target && d.target.data && d.target.data.isDashed);
    return `${sourceId} -> ${targetId} (Dashed: ${isDashed})`;
  }));


  // 1. Update the Nodes
  const node = gChart.selectAll('g.node')
      .data(nodesToRender, d => d.id); // Use nodesToRender here

  const nodeEnter = node.enter().append('g')
      .attr('class', 'node')
      .attr('transform', d => `translate(${source.x0},${source.y0})`)
      .on('click', click);

  nodeEnter.append('rect')
      .attr('width', 0)
      .attr('height', 0)
      .attr('rx', 8)
      .attr('ry', 8)
      .style('fill', d => getNodeColor(d.data.id))
      .attr('stroke', d => d.data.isDashed ? '#888' : '#ccc') // Node's own dashed status
      .attr('stroke-width', 1.5)
      .attr('stroke-dasharray', d => d.data.isDashed ? ("5, 5") : "none"); // Node's own dashed status

  nodeEnter.append('text')
      .attr('dy', '.35em')
      .attr('x', 0)
      .attr('y', 0)
      .attr('text-anchor', 'middle')
      .style('fill', d => getTextColor(d.data.id))
      .style('font-size', '14px')
      .text(d => appStore.checkLang.isTh ? d.data.name_th : d.data.name_en)
      .style('opacity', 0);

  const nodeUpdate = nodeEnter.merge(node);

  const transitionsToWaitFor = [];

  const nodeTransition = nodeUpdate.transition()
      .duration(750)
      .attr('transform', d => `translate(${d.x},${d.y})`);
  transitionsToWaitFor.push(endAll(nodeTransition));

  const textTransition = nodeUpdate.select('text')
      .transition()
      .duration(750)
      .style('opacity', 1)
      .tween("text", function(d) {
        const textElement = d3.select(this);
        return function() {
          const bbox = textElement.node().getBBox();
          const paddingX = 20;
          const paddingY = 15;
          let rectWidth = bbox.width + paddingX * 2;
          let rectHeight = bbox.height + paddingY * 2;

          rectWidth = Math.max(rectWidth, BASE_NODE_WIDTH);
          rectHeight = Math.max(rectHeight, BASE_NODE_HEIGHT);

          d.data.rectWidth = rectWidth;
          d.data.rectHeight = rectHeight;

          d3.select(this.previousElementSibling)
              .attr('width', rectWidth)
              .attr('height', rectHeight)
              .attr('x', -rectWidth / 2)
              .attr('y', -rectHeight / 2);
        };
      });
  transitionsToWaitFor.push(endAll(textTransition));

  nodeUpdate.select('rect') // Ensure stroke/dasharray is updated for existing nodes too
      .attr('stroke', d => d.data.isDashed ? '#888' : '#ccc')
      .attr('stroke-dasharray', d => d.data.isDashed ? ("5, 5") : "none");

  const nodeExit = node.exit().transition()
      .duration(750)
      .attr('transform', d => `translate(${source.x},${source.y})`)
      .remove();
  transitionsToWaitFor.push(endAll(nodeExit));

  nodeExit.select('rect')
      .attr('width', 0)
      .attr('height', 0);

  nodeExit.select('text')
      .style('opacity', 0);


  // 2. Update the Links
  const link = gChart.selectAll('path.link')
      // Use a robust key, checking for existence of source and target IDs
      .data(links, d => `${d.source ? d.source.id : 'null'}-${d.target ? d.target.id : 'null'}`);

  const linkEnter = link.enter().insert('path', 'g') // Insert before 'g' nodes
      .attr('class', 'link')
      .attr('fill', 'none')
      .attr('stroke', '#ccc')
      .attr('stroke-width', 1.5)
      // !!! CRITICAL CHANGE HERE for linkEnter stroke-dasharray !!!
      // Check if link.data.isDashed (for custom links) or link.target.data.isDashed (for standard links)
      .attr('stroke-dasharray', d => (d.data && d.data.isDashed) || (d.target && d.target.data && d.target.data.isDashed) ? ("5, 5") : "none")
      .attr('d', d => {
        // Determine the animation source based on the actual source of the link
        const animationSource = (isArrayRoot && d.source.id === 'dummy_root' && mainRoot) ? mainRoot : d.source;

        // Ensure animationSource has the required properties
        const startNodeForEnter = {
          x: animationSource.x0,
          y: animationSource.y0,
          data: {
            rectWidth: animationSource.data?.rectWidth || BASE_NODE_WIDTH,
            rectHeight: animationSource.data?.rectHeight || BASE_NODE_HEIGHT
          }
        };
        return diagonal(startNodeForEnter, startNodeForEnter);
      });

  const linkUpdate = linkEnter.merge(link);

  const linkTransition = linkUpdate.transition()
      .duration(750)
      .attr('d', d => diagonal(d.source, d.target));
  transitionsToWaitFor.push(endAll(linkTransition));

  // !!! CRITICAL CHANGE HERE for linkUpdate stroke-dasharray (LINE 318) !!!
  // Ensure stroke-dasharray is updated for existing links as well
  linkUpdate.attr('stroke-dasharray', d => (d.data && d.data.isDashed) || (d.target && d.target.data && d.target.data.isDashed) ? ("5, 5") : "none");


  const linkExit = link.exit().transition()
      .duration(750)
      .attr('d', d => {
        // Determine the animation source based on the actual source of the link
        const animationSource = (isArrayRoot && d.source.id === 'dummy_root' && mainRoot) ? mainRoot : d.source;

        // Ensure animationSource has the required properties
        const endNodeForExit = {
          x: animationSource.x,
          y: animationSource.y,
          data: {
            rectWidth: animationSource.data?.rectWidth || BASE_NODE_WIDTH,
            rectHeight: animationSource.data?.rectHeight || BASE_NODE_HEIGHT
          }
        };
        return diagonal(endNodeForExit, endNodeForExit);
      })
      .remove();
  transitionsToWaitFor.push(endAll(linkExit));

  // Stash the current positions (x, y) as old positions (x0, y0) for the next transition.
  // Only stash for visible nodes.
  visibleNodes.forEach(d => {
    d.x0 = d.x;
    d.y0 = d.y;
  });

  // *** IMPORTANT CHANGE: Wait for all transitions to complete before fitting to view ***
  Promise.all(transitionsToWaitFor)
      .then(() => {
        console.log("All transitions ended. Calling fitToView.");
        fitToView();
      })
      .catch(error => {
        console.error("Error during D3 transitions:", error);
      });
};

function click(event, d) {
  console.log(`Click event on node: ${d.id} (${d.data.name_th || d.data.name_en})`);

  // Handle clicking on the dummy root for array data (if applicable)
  const isArrayRoot = Array.isArray(props.treeData);
  if (isArrayRoot && d.id === 'dummy_root') {
    // Optionally, prevent clicking the dummy root or handle it specifically
    console.log("Clicked dummy root, no action.");
    return;
  }

  if (d.children) {
    d._children = d.children; // Store current children
    d.children = null;       // Hide children
  }
  else if (d._children) {
    d.children = d._children; // Restore children
    d._children = null;       // Clear hidden children
  }
  else {
    console.log(`  Leaf node or no hidden children. No action taken.`);
  }
  update(d); // Trigger update with the clicked node as the source for animation
}

// DIAGONAL FUNCTION (unchanged)
function diagonal(s, d) {
  // Ensure s.data and d.data exist before accessing rectWidth/Height
  const sourceRectWidth = s.data?.rectWidth || BASE_NODE_WIDTH;
  const sourceRectHeight = s.data?.rectHeight || BASE_NODE_HEIGHT;
  const targetRectWidth = d.data?.rectWidth || BASE_NODE_WIDTH;
  const targetRectHeight = d.data?.rectHeight || BASE_NODE_HEIGHT;

  const startX = s.x;
  const startY = s.y + sourceRectHeight / 2; // Link starts from bottom center of source node

  const endX = d.x;
  const endY = d.y - targetRectHeight / 2; // Link ends at top center of target node

  const midY = startY + (endY - startY) / 2; // Midpoint for the vertical line segment

  return `M ${startX},${startY}
          V ${midY}
          H ${endX}
          V ${endY}`;
}

const fitToView = () => {
  if (!gChart || gChart.empty()) {
    console.warn("gChart not ready or empty for fitToView, retrying...");
    setTimeout(fitToView, 500); // Short retry if gChart is not ready
    return;
  }

  const treeBBox = gChart.node()?.getBBox();

  if (!treeBBox || treeBBox.width === 0 || treeBBox.height === 0) {
    console.warn("Tree BBox is zero or invalid, retrying fit...");
    setTimeout(fitToView, 500); // Reduce delay for retry
    return;
  }

  const containerWidth = diagramContainer.value.offsetWidth;
  const containerHeight = diagramContainer.value.offsetHeight;

  const treeWidth = treeBBox.width;
  const treeHeight = treeBBox.height;

  const buffer = 0.90; // You can adjust this value (e.g., 0.85, 0.8)

  const scaleX = (containerWidth * buffer) / treeWidth;
  const scaleY = (containerHeight * buffer) / treeHeight;
  const scale = Math.min(scaleX, scaleY);

  const translateX = (containerWidth / 2) - ((treeBBox.x + treeWidth / 2) * scale);
  const translateY = (containerHeight / 2) - ((treeBBox.y + treeHeight / 2) * scale);

  console.log("Calculated BBox:", { x: treeBBox.x, y: treeBBox.y, width: treeWidth, height: treeHeight });
  console.log("Container Size:", { containerWidth, containerHeight });
  console.log("Calculated Scale:", scale);
  console.log("Calculated Translate:", { translateX, translateY });

  gChart.transition()
      .duration(750)
      .attr("transform", `translate(${translateX},${translateY})scale(${scale})`);
};

onMounted(() => {
  setupD3Tree();
  window.addEventListener('resize', setupD3Tree); // Call setupD3Tree on resize to re-render
});

watch(() => props.treeData, (newVal, oldVal) => {
  // Deep watch to ensure changes in nested properties trigger re-render
  if (JSON.stringify(newVal) !== JSON.stringify(oldVal)) {
    console.log("props.treeData content changed, re-setting up D3 Tree.");
    setupD3Tree(false); // Do not preserve state on new data
  } else {
    console.log("props.treeData reference changed, but content is same. Ignoring watch for internal D3 state changes.");
  }
}, { deep: true });

watch(() => appStore.checkLang.isTh, () => {
  console.log("Language changed, re-setting up D3 Tree.");
  setupD3Tree(false); // Re-render on language change
});

onUnmounted(() => {
  window.removeEventListener('resize', setupD3Tree);
});
</script>

<style scoped>
.diagram-container {
  width: 100%;
  height: 90vh;
  background-color: white;
  overflow: hidden; /* Important for containing the SVG */
  display: flex;
  justify-content: center; /* Center SVG horizontally if it's smaller than container */
  align-items: center;   /* Center SVG vertically if it's smaller than container */
}

/* Ensure SVG fills the container */
.diagram-container svg {
  display: block; /* Remove extra space below SVG */
}

.node rect {
  cursor: pointer;
}

.node text {
  font-family: 'Arial', sans-serif;
  cursor: pointer;
  user-select: none;
}

.link {
  /* Stroke and fill are set dynamically in JS */
}
</style>