import ReactChart from '../../../models/ReactChart';
import { ChartConfig } from 'app/types/ChartConfig';
import Config from './config';
import VChart from '@visactor/vchart';
import { SankeyStyleConfig } from './types';
import React, { useEffect, useRef, useCallback } from 'react';
import { BrokerContext } from 'app/types/ChartLifecycleBroker';

// 添加节点类型定义
interface SankeyNode {
  name: string;
  value: number;
  children?: Map<string, SankeyNode>;
}

const SankeyChartComponent = React.memo((props: any) => {
  const { dataset, style, setChart } = props;
  const sankeyStyle = style as SankeyStyleConfig;
  const chartRef = useRef<HTMLDivElement>(null);
  const chartInstance = useRef<any>(null);

  const getChartSpec = useCallback(() => {
    const defaultSpec = {
      type: 'sankey',
      background: 'transparent',
      data: [
        {
          name: 'data',
          values: [
            {
              nodes: [
                  {
                    value: 100,
                    name: 'A',
                    children: [
                      {
                        name: 'top',
                        value: 40,
                        children: [
                          { name: '00', value: 15 },
                          { name: '01', value: 10 },
                          { name: '02', value: 10 }
                        ]
                      },
                      {
                        name: 'middle',
                        value: 30,
                        children: [
                          { name: '00', value: 10 },
                          { name: '01', value: 10 },
                          { name: '02', value: 10 }
                        ]
                      },
                      {
                        name: 'bottom',
                        value: 30
                      }
                    ]
                  },
                  {
                    value: 80,
                    name: 'B',
                    children: [
                      {
                        name: 'top',
                        value: 40,
                        children: [
                          { name: '00', value: 100 },
                          { name: '01', value: 40 }
                        ]
                      },
                      {
                        name: 'middle',
                        value: 10
                      },
                      {
                        name: 'bottom',
                        value: 30
                      }
                    ]
                  },
                  {
                    value: 50,
                    name: 'C',
                    children: [
                      {
                        name: 'top',
                        value: 20
                      },
                      {
                        name: 'middle',
                        value: 20
                      },
                      {
                        name: 'bottom',
                        value: 10
                      }
                    ]
                  }
                ]
            }
          ]
        }
      ],
      categoryField: 'name',
      valueField: 'value',
      nodeAlign: sankeyStyle?.sankeyTitle?.nodeAlign || 'left',
      nodeGap: sankeyStyle?.sankeyTitle?.nodeGap || 8,
      nodeWidth: sankeyStyle?.sankeyTitle?.nodeWidth || 10,
      minNodeHeight: 4,
      nodeKey: (datum: any) => datum.name,

      label: {
        visible: sankeyStyle?.label?.showLabel ?? true,
        style: {
          fill: '#333333',
          fontSize: 12
        },
        state: {
          blur: {
            fill: '#e8e8e8',
            fillOpacity: 0.15
          }
        }
      },

      node: {
        style: {
          fill: '#B4E1FF',
          fillOpacity: 0.7
        },
        state: {
          hover: {
            fill: sankeyStyle?.nodeStyle?.hoverColor || 'red'
          },
          blur: {
            fill: '#e8e8e8',
            fillOpacity: 0.15
          }
        }
      },
      link: {
        state: {
          hover: {
            stroke: 'yellow'
          },
          blur: {
            fill: '#e8e8e8'
          }
        }
      },

      emphasis: {
        enable: true,
        effect: 'related'
      }
    };

    if (!dataset?.rows?.length || !dataset?.columns?.length) {
      return defaultSpec;
    }

    try {
      const dimensionFields = (dataset.fields?.dimension || [])
        .map(field => field.colName);
      
      if (dimensionFields.length < 2) {
        console.warn('Sankey chart requires at least 2 dimensions');
        return defaultSpec;
      }

      const dimensionIndices = dimensionFields.map(fieldName => 
        dataset.columns.findIndex(col => col.name[0] === fieldName)
      );

      const valueIndex = dataset.columns.findIndex(
        col => col.name[0] === dataset.fields?.metrics?.[0]?.colName
      );

      if (dimensionIndices.includes(-1) || valueIndex === -1) {
        console.warn('Cannot find required fields in columns');
        return defaultSpec;
      }

      const hierarchyMap = new Map<string, SankeyNode>();

      dataset.rows.forEach(row => {
        const firstDimValue = String(row[dimensionIndices[0]]) || 'Unknown';
        const secondDimValue = String(row[dimensionIndices[1]]) || 'Unknown';
        const thirdDimValue = dimensionIndices[2] ? String(row[dimensionIndices[2]]) || 'Unknown' : null;
        const value = Number(row[valueIndex] || 0);

        if (!hierarchyMap.has(firstDimValue)) {
          hierarchyMap.set(firstDimValue, {
            name: firstDimValue,
            value: 0,
            children: new Map()
          });
        }

        const firstNode = hierarchyMap.get(firstDimValue)!;
        firstNode.value += value;

        if (!firstNode.children!.has(secondDimValue)) {
          firstNode.children!.set(secondDimValue, {
            name: secondDimValue,
            value: 0,
            children: thirdDimValue ? new Map() : undefined
          });
        }

        const secondNode = firstNode.children!.get(secondDimValue)!;
        secondNode.value += value;

        if (thirdDimValue && secondNode.children) {
          if (!secondNode.children.has(thirdDimValue)) {
            secondNode.children.set(thirdDimValue, {
              name: thirdDimValue,
              value: 0
            });
          }

          const thirdNode = secondNode.children.get(thirdDimValue)!;
          thirdNode.value += value;
        }
      });

      const nodes = Array.from(hierarchyMap.values()).map(node => ({
        name: node.name,
        value: node.value,
        children: Array.from(node.children!.values()).map(child => ({
          name: child.name,
          value: child.value,
          children: child.children ? Array.from(child.children.values()) : undefined
        }))
      }));

      return {
        ...defaultSpec,
        data: [
          {
            name: 'data',
            values: [{ nodes }]
          }
        ]
      };
    } catch (error) {
      console.error('Failed to process sankey chart data:', error);
      return defaultSpec;
    }
  }, [dataset, sankeyStyle]);

  const onUpdate = useCallback(() => {
    if (!chartInstance.current || !chartRef.current) return;

    try {
      const spec = getChartSpec();
      chartInstance.current.updateSpec(spec);
      chartInstance.current.renderSync();
    } catch (error) {
      console.error('Failed to update sankey chart:', error);
    }
  }, [getChartSpec]);

  useEffect(() => {
    if (!chartRef.current) return;

    if (chartInstance.current) {
      chartInstance.current.release();
    }

    const spec = getChartSpec();
    chartInstance.current = new VChart(spec, { dom: chartRef.current });
    chartInstance.current.renderSync();
    
    setChart?.(chartInstance.current);

    return () => {
      if (chartInstance.current) {
        chartInstance.current.release();
        setChart?.(null);
      }
    };
  }, [getChartSpec, setChart]);

  useEffect(() => {
    if (chartInstance.current) {
      onUpdate();
    }
  }, [dataset, sankeyStyle, onUpdate]);

  return <div ref={chartRef} style={{ width: '100%', height: '100%' }} />;
});

class SankeyChartClass extends ReactChart {
  public chart: any;
  public config: ChartConfig;

  constructor() {
    super(SankeyChartComponent, {
      id: 'sankey',
      name: 'viz.chart.sankey',
      icon: 'sankey',
    });
    this.meta.requirements = [
      {
        group: 2,
        aggregate: 1,
      },
    ];
    this.chart = null;
    this.config = Config;
  }

  public onMount(options: any, context: BrokerContext): void {
    console.log('options', options);
    if (options.containerId === undefined || !context.document) {
      return;
    }

    const spec = this.getChartSpec(options.dataset, options.config.styles);
    console.log('spec', spec);

    this.chart = new VChart(spec, { 
      dom: context.document.getElementById(options.containerId)!
    });
    this.chart.renderSync();
  }

  public onUpdated(options: any, context?: any): void {
    console.log('onUpdated', options);
    if (!options.dataset || !options.dataset.columns || !options.config) {
      return;
    }

    if (!this.isMatchRequirement(options.config)) {
      this.chart?.clear();
      return;
    }

    try {
      const spec = this.getChartSpec(options.dataset, options.config.styles);
      this.chart?.updateSpec(spec);
      this.chart?.renderSync();
    } catch (error) {
      console.error('Failed to update sankey chart:', error);
    }
  }

  public onUnMount(options: any, context: BrokerContext): void {
    if (this.chart) {
      this.chart.release();
      this.chart = null;
    }
  }

  private transformData(dataset: any) {
    if (!dataset?.rows?.length || !dataset?.columns?.length) {
      return [{ nodes: [] }];
    }

    try {
      const dimensionFields = (dataset.columns || [])
        .filter(field => field.type === 'STRING')
        .map(field => field.name[0]);

      const valueFields = (dataset.columns || [])
        .filter(field => field.type === 'NUMERIC')
        .map(field => field.name[0]);
      
      if (dimensionFields.length < 2 || valueFields.length < 1) {
        return [{ nodes: [] }];
      }

      const dimensionIndices = dimensionFields.map(fieldName => 
        dataset.columns.findIndex(col => col.name[0] === fieldName)
      );

      const valueIndex = valueFields.map(fieldName => 
        dataset.columns.findIndex(col => col.name[0] === fieldName)
      )[0];

      if (dimensionIndices.includes(-1) || valueIndex === -1) {
        return [{ nodes: [] }];
      }

      const hierarchyMap = new Map<string, SankeyNode>();

      dataset.rows.forEach(row => {
        const firstDimValue = String(row[dimensionIndices[0]]) || 'Unknown';
        const secondDimValue = String(row[dimensionIndices[1]]) || 'Unknown';
        const thirdDimValue = dimensionIndices[2] ? String(row[dimensionIndices[2]]) || 'Unknown' : null;
        const value = Number(row[valueIndex] || 0);

        if (!hierarchyMap.has(firstDimValue)) {
          hierarchyMap.set(firstDimValue, {
            name: firstDimValue,
            value: 0,
            children: new Map()
          });
        }

        const firstNode = hierarchyMap.get(firstDimValue)!;
        firstNode.value += value;

        if (!firstNode.children!.has(secondDimValue)) {
          firstNode.children!.set(secondDimValue, {
            name: secondDimValue,
            value: 0,
            children: thirdDimValue ? new Map() : undefined
          });
        }

        const secondNode = firstNode.children!.get(secondDimValue)!;
        secondNode.value += value;

        if (thirdDimValue && secondNode.children) {
          if (!secondNode.children.has(thirdDimValue)) {
            secondNode.children.set(thirdDimValue, {
              name: thirdDimValue,
              value: 0
            });
          }

          const thirdNode = secondNode.children.get(thirdDimValue)!;
          thirdNode.value += value;
        }
      });

      const nodes = Array.from(hierarchyMap.values()).map(node => ({
        name: node.name,
        value: node.value,
        children: Array.from(node.children!.values()).map(child => ({
          name: child.name,
          value: child.value,
          children: child.children ? Array.from(child.children.values()) : undefined
        }))
      }));

      return [{ nodes }];
    } catch (error) {
      console.error('Failed to transform data:', error);
      return [{ nodes: [] }];
    }
  }

  private getChartSpec(dataset: any, styles: any) {
    const defaultSpec = {
      type: 'sankey',
      background: 'transparent',
      data: [
        {
          name: 'data',
          values: [{ nodes: [] }]
        }
      ],
      categoryField: 'name',
      valueField: 'value',
      nodeAlign: styles?.find(s => s.key === 'sankeyTitle')?.rows?.find(r => r.key === 'nodeAlign')?.value || 'left',
      nodeGap: styles?.find(s => s.key === 'sankeyTitle')?.rows?.find(r => r.key === 'nodeGap')?.value || 8,
      nodeWidth: styles?.find(s => s.key === 'sankeyTitle')?.rows?.find(r => r.key === 'nodeWidth')?.value || 10,
      minNodeHeight: 4,
      nodeKey: (datum: any) => datum.name,

      label: {
        visible: styles?.find(s => s.key === 'label')?.rows?.find(r => r.key === 'showLabel')?.value ?? true,
        style: {
          fill: '#333333',
          fontSize: 12
        }
      },

      node: {
        style: {
          fill: '#B4E1FF',
          fillOpacity: 0.7
        },
        state: {
          hover: {
            fill: styles?.find(s => s.key === 'nodeStyle')?.rows?.find(r => r.key === 'hoverColor')?.value || '#1890ff'
          },
          blur: {
            fill: '#e8e8e8',
            fillOpacity: 0.15
          }
        }
      },

      link: {
        state: {
          hover: {
            stroke: '#000000'
          },
          blur: {
            fill: '#e8e8e8'
          }
        }
      },

      emphasis: {
        enable: true,
        effect: 'adjacency'
      }
    };

    if (!dataset?.rows?.length || !dataset?.columns?.length) {
      console.log('defaultSpec', defaultSpec);
      return defaultSpec;
    }

    return {
      ...defaultSpec,
      data: [
        {
          name: 'data',
          values: this.transformData(dataset)
        }
      ]
    };
  }
}

export default SankeyChartClass;