candle-annotator/src/components/CandleChart.tsx

'use client';

import { useEffect, useRef, useState, useImperativeHandle, forwardRef } from 'react';
import { createChart, IChartApi, ISeriesApi, CandlestickData, Time } from 'lightweight-charts';
import SvgOverlay from './SvgOverlay';

interface Candle {
  time: number;
  open: number;
  high: number;
  low: number;
  close: number;
}

interface Annotation {
  id: number;
  timestamp: number;
  label_type: string;
  geometry: {
    startTime?: number;
    startPrice?: number;
    endTime?: number;
    endPrice?: number;
  } | null;
  created_at: number;
}

interface CandleChartProps {
  activeTool: string | null;
  onAnnotationChange?: () => void;
}

export interface CandleChartHandle {
  refreshData: () => Promise<void>;
}

const CandleChart = forwardRef<CandleChartHandle, CandleChartProps>(
  ({ activeTool, onAnnotationChange }, ref) => {
    const chartContainerRef = useRef<HTMLDivElement>(null);
    const chartRef = useRef<IChartApi | null>(null);
    const seriesRef = useRef<ISeriesApi<'Candlestick'> | null>(null);
    const [candles, setCandles] = useState<Candle[]>([]);
    const [annotations, setAnnotations] = useState<Annotation[]>([]);
    const [isEmpty, setIsEmpty] = useState(true);

    // Fetch candles from API
    const fetchCandles = async () => {
      try {
        const response = await fetch('/api/candles');
        const data = await response.json();
        setCandles(data);
        setIsEmpty(data.length === 0);
        return data;
      } catch (error) {
        console.error('Failed to fetch candles:', error);
        return [];
      }
    };

    // Fetch annotations from API
    const fetchAnnotations = async () => {
      try {
        const response = await fetch('/api/annotations');
        const data = await response.json();
        setAnnotations(data);
        return data;
      } catch (error) {
        console.error('Failed to fetch annotations:', error);
        return [];
      }
    };

    // Expose refresh method to parent
    useImperativeHandle(ref, () => ({
      refreshData: async () => {
        await fetchCandles();
        await fetchAnnotations();
      },
    }));

    // Initialize chart
    useEffect(() => {
      if (!chartContainerRef.current || isEmpty) return;

      const chart = createChart(chartContainerRef.current, {
        width: chartContainerRef.current.clientWidth,
        height: chartContainerRef.current.clientHeight,
        layout: {
          background: { color: '#0f172a' }, // Slate-900
          textColor: '#e2e8f0', // Slate-200
        },
        grid: {
          vertLines: { color: '#1e293b' }, // Subtle grid
          horzLines: { color: '#1e293b' },
        },
        crosshair: {
          mode: 1,
        },
        timeScale: {
          timeVisible: true,
          secondsVisible: false,
          borderColor: '#334155',
        },
        rightPriceScale: {
          borderColor: '#334155',
        },
      });

      const candlestickSeries = chart.addCandlestickSeries({
        upColor: '#22c55e', // Green
        downColor: '#ef4444', // Red
        borderVisible: false,
        wickUpColor: '#22c55e',
        wickDownColor: '#ef4444',
      });

      chartRef.current = chart;
      seriesRef.current = candlestickSeries;

      // Handle resize
      const resizeObserver = new ResizeObserver((entries) => {
        if (entries.length === 0 || !chartContainerRef.current) return;
        const { width, height } = entries[0].contentRect;
        chart.applyOptions({ width, height });
      });

      resizeObserver.observe(chartContainerRef.current);

      return () => {
        resizeObserver.disconnect();
        chart.remove();
      };
    }, [isEmpty]);

    // Load candle data into chart
    useEffect(() => {
      if (!seriesRef.current || candles.length === 0) return;

      const chartData: CandlestickData[] = candles
        .map((candle) => ({
          time: candle.time as Time,
          open: candle.open,
          high: candle.high,
          low: candle.low,
          close: candle.close,
        }))
        .sort((a, b) => (a.time as number) - (b.time as number));

      seriesRef.current.setData(chartData);
    }, [candles]);

    // Update markers from annotations
    useEffect(() => {
      if (!seriesRef.current) return;

      const markerAnnotations = annotations.filter(
        (a) => a.label_type === 'break_up' || a.label_type === 'break_down'
      );

      const markers = markerAnnotations
        .map((annotation) => ({
          time: annotation.timestamp as Time,
          position: annotation.label_type === 'break_up' ? ('belowBar' as const) : ('aboveBar' as const),
          color: annotation.label_type === 'break_up' ? '#22c55e' : '#ef4444',
          shape: annotation.label_type === 'break_up' ? ('arrowUp' as const) : ('arrowDown' as const),
          text: annotation.label_type === 'break_up' ? 'Break Up' : 'Break Down',
        }))
        .sort((a, b) => (a.time as number) - (b.time as number));

      seriesRef.current.setMarkers(markers);
    }, [annotations]);

    // Handle chart clicks for annotation
    useEffect(() => {
      if (!chartRef.current || !seriesRef.current) return;

      const handleClick = async (param: any) => {
        if (!param.point || !activeTool) return;

        const timeCoordinate = param.point.x;
        const priceCoordinate = param.point.y;

        const time = chartRef.current!.timeScale().coordinateToTime(timeCoordinate);
        const price = seriesRef.current!.coordinateToPrice(priceCoordinate);

        if (time === null || price === null) return;

        // For break_up and break_down, snap to nearest candle
        if (activeTool === 'break_up' || activeTool === 'break_down') {
          const timestamp = typeof time === 'string' ? Date.parse(time) / 1000 : (time as number);
          
          // Find nearest candle
          const nearestCandle = candles.reduce((prev, curr) => {
            return Math.abs(curr.time - timestamp) < Math.abs(prev.time - timestamp) ? curr : prev;
          }, candles[0]);

          if (!nearestCandle) return;

          // Create annotation
          try {
            const response = await fetch('/api/annotations', {
              method: 'POST',
              headers: { 'Content-Type': 'application/json' },
              body: JSON.stringify({
                timestamp: nearestCandle.time,
                label_type: activeTool,
              }),
            });

            if (response.ok) {
              await fetchAnnotations();
              onAnnotationChange?.();
            }
          } catch (error) {
            console.error('Failed to create annotation:', error);
          }
        }

        // For delete tool, find and delete marker at clicked position
        if (activeTool === 'delete') {
          const timestamp = typeof time === 'string' ? Date.parse(time) / 1000 : (time as number);
          
          // Find annotation at this timestamp (within tolerance)
          const tolerance = 60; // 60 seconds tolerance
          const annotation = annotations.find(
            (a) =>
              (a.label_type === 'break_up' || a.label_type === 'break_down') &&
              Math.abs(a.timestamp - timestamp) < tolerance
          );

          if (annotation) {
            try {
              const response = await fetch(`/api/annotations/${annotation.id}`, {
                method: 'DELETE',
              });

              if (response.ok) {
                await fetchAnnotations();
                onAnnotationChange?.();
              }
            } catch (error) {
              console.error('Failed to delete annotation:', error);
            }
          }
        }
      };

      chartRef.current.subscribeClick(handleClick);

      return () => {
        chartRef.current?.unsubscribeClick(handleClick);
      };
    }, [activeTool, candles, annotations, onAnnotationChange]);

    // Fetch data on mount
    useEffect(() => {
      fetchCandles();
      fetchAnnotations();
    }, []);

    if (isEmpty) {
      return (
        <div className="flex-1 flex items-center justify-center bg-background">
          <div className="text-muted-foreground text-center">
            <p className="text-lg">Upload a CSV file to view the candlestick chart</p>
            <p className="text-sm mt-2">CSV format: time, open, high, low, close</p>
          </div>
        </div>
      );
    }

    return (
      <div className="w-full h-full relative">
        <div ref={chartContainerRef} className="absolute inset-0" />
        <SvgOverlay
          chart={chartRef.current}
          series={seriesRef.current}
          activeTool={activeTool}
          onAnnotationChange={onAnnotationChange}
        />
      </div>
    );
  }
);

CandleChart.displayName = 'CandleChart';

export default CandleChart;