candle-annotator/src/app/page.tsx

'use client';

import { useState, useRef, useEffect, useCallback } from 'react';
import Toolbox, { Tool } from '@/components/Toolbox';
import FileUpload from '@/components/FileUpload';
import CandleChart, { CandleChartHandle } from '@/components/CandleChart';
import ChartSelector from '@/components/ChartSelector';
import PredictionPanel from '@/components/PredictionPanel';
import type { PredictionState, PredictionSpan, ModelInfoResponse, Disagreement, DisagreementType, PredictionSummary } from '@/types/predictions';

/**
 * Calculate overlap between two time ranges
 * Returns overlap ratio (0-1) relative to the smaller span
 */
function calculateOverlap(
  start1: number,
  end1: number,
  start2: number,
  end2: number
): number {
  const overlapStart = Math.max(start1, start2);
  const overlapEnd = Math.min(end1, end2);
  const overlap = Math.max(0, overlapEnd - overlapStart);
  const minLength = Math.min(end1 - start1, end2 - start2);
  return minLength > 0 ? overlap / minLength : 0;
}

/**
 * Detect disagreements between human annotations and model predictions
 * Returns summary with agreement/disagreement counts and detailed disagreement list
 */
function detectDisagreements(
  humanSpans: SpanAnnotation[],
  predictionSpans: PredictionSpan[],
  overlapThreshold: number = 0.5
): PredictionSummary {
  const disagreements: Disagreement[] = [];
  const matchedHumanIds = new Set<number>();
  const matchedPredictionIndices = new Set<number>();
  
  // Filter human spans to only those with a source of "human" (not programmatic)
  const humanOnlySpans = humanSpans.filter(span => {
    // Assuming human-created spans don't have a source field or have source="human"
    // This will be properly implemented when we add the source field
    return true; // For now, consider all spans
  });

  // Find matches and label mismatches
  humanOnlySpans.forEach((humanSpan) => {
    predictionSpans.forEach((predSpan, predIdx) => {
      const overlap = calculateOverlap(
        humanSpan.start_time,
        humanSpan.end_time,
        predSpan.start_time,
        predSpan.end_time
      );

      if (overlap >= overlapThreshold) {
        matchedHumanIds.add(humanSpan.id);
        matchedPredictionIndices.add(predIdx);

        // Check for label mismatch
        if (humanSpan.label !== predSpan.label) {
          disagreements.push({
            type: 'label_mismatch',
            humanSpan: {
              id: humanSpan.id,
              label: humanSpan.label,
              start_time: humanSpan.start_time,
              end_time: humanSpan.end_time,
            },
            predictionSpan: predSpan,
            overlap_ratio: overlap,
          });
        }
      }
    });
  });

  // Find spans missed by model (human annotation but no prediction)
  humanOnlySpans.forEach((humanSpan) => {
    if (!matchedHumanIds.has(humanSpan.id)) {
      disagreements.push({
        type: 'missed_by_model',
        humanSpan: {
          id: humanSpan.id,
          label: humanSpan.label,
          start_time: humanSpan.start_time,
          end_time: humanSpan.end_time,
        },
      });
    }
  });

  // Find spans missed by human (prediction but no human annotation)
  predictionSpans.forEach((predSpan, idx) => {
    if (!matchedPredictionIndices.has(idx)) {
      disagreements.push({
        type: 'missed_by_human',
        predictionSpan: predSpan,
      });
    }
  });

  // Calculate agreements (matched spans with same label)
  const agreements = matchedHumanIds.size - disagreements.filter(d => d.type === 'label_mismatch').length;

  return {
    total_predictions: predictionSpans.length,
    total_human_annotations: humanOnlySpans.length,
    agreements,
    disagreements,
  };
}

interface Chart {
  id: number;
  name: string;
  created_at: number;
}

interface Annotation {
  id: number;
  chart_id: number;
  timestamp: number;
  label_type: string;
  geometry: any;
  created_at: number;
}

interface SpanAnnotation {
  id: number;
  chart_id: number;
  start_time: number;
  end_time: number;
  label: string;
  confidence: number | null;
  outcome: string | null;
  notes: string | null;
  sub_spans: any;
  color: string;
  created_at: number;
}

interface SpanLabelType {
  id: number;
  name: string;
  display_name: string;
  color: string;
  hotkey: string | null;
  is_active: number;
  sort_order: number;
  created_at: number;
}

export default function Home() {
  const [activeTool, setActiveTool] = useState<Tool | 'span'>(null);
  const [selectedColor, setSelectedColor] = useState('#3b82f6');
  const [selectedLabelId, setSelectedLabelId] = useState<number | null>(null);
  const [annotations, setAnnotations] = useState<Annotation[]>([]);
  const [charts, setCharts] = useState<Chart[]>([]);
  const [activeChartId, setActiveChartId] = useState<number | null>(null);
  const chartRef = useRef<CandleChartHandle>(null);

  // Span annotation state
  const [spanAnnotations, setSpanAnnotations] = useState<SpanAnnotation[]>([]);
  const [selectedSpanId, setSelectedSpanId] = useState<number | null>(null);
  const [spanLabelTypes, setSpanLabelTypes] = useState<SpanLabelType[]>([]);

  // Prediction state
  const [predictionState, setPredictionState] = useState<PredictionState>({
    spans: [],
    perCandlePredictions: [],
    isLoading: false,
    error: null,
    modelInfo: null,
    visible: false,
    confidenceThreshold: 0.5,
    selectedLabels: new Set<string>(),
    autoPredict: false,
    cacheKey: null,
  });

  // Prediction cache: Map<cacheKey, { spans, predictions, modelVersion }>
  const predictionCacheRef = useRef<Map<string, {
    spans: PredictionSpan[];
    predictions: any[];
    modelVersion: string;
  }>>(new Map());

  // Model health state
  const [isModelOnline, setIsModelOnline] = useState(true);

  // Prediction summary state
  const [predictionSummary, setPredictionSummary] = useState<PredictionSummary | null>(null);
  
  // Disagreement filter state
  const [showOnlyDisagreements, setShowOnlyDisagreements] = useState(false);

  // Fetch charts list
  const fetchCharts = useCallback(async () => {
    try {
      const response = await fetch('/api/charts');
      const data = await response.json();
      setCharts(data);
      return data as Chart[];
    } catch (error) {
      console.error('Failed to fetch charts:', error);
      return [];
    }
  }, []);

  // Fetch annotations for active chart
  const fetchAnnotations = useCallback(async (chartId: number | null) => {
    if (!chartId) {
      setAnnotations([]);
      return;
    }
    try {
      const response = await fetch(`/api/annotations?chartId=${chartId}`);
      const data = await response.json();
      setAnnotations(data);
    } catch (error) {
      console.error('Failed to fetch annotations:', error);
    }
  }, []);

  // Fetch span annotations for active chart
  const fetchSpanAnnotations = useCallback(async (chartId: number | null) => {
    if (!chartId) {
      setSpanAnnotations([]);
      return;
    }
    try {
      const response = await fetch(`/api/span-annotations?chartId=${chartId}`);
      const data = await response.json();
      setSpanAnnotations(data);
    } catch (error) {
      console.error('Failed to fetch span annotations:', error);
    }
  }, []);

  // Fetch span label types
  const fetchSpanLabelTypes = useCallback(async () => {
    try {
      const response = await fetch('/api/span-label-types');
      const data = await response.json();
      setSpanLabelTypes(data);
    } catch (error) {
      console.error('Failed to fetch span label types:', error);
    }
  }, []);

  // Fetch charts and span label types on mount, auto-select the most recent chart
  useEffect(() => {
    const init = async () => {
      const chartList = await fetchCharts();
      await fetchSpanLabelTypes();
      if (chartList.length > 0) {
        setActiveChartId(chartList[0].id); // sorted by created_at desc
      }
    };
    init();
  }, [fetchCharts, fetchSpanLabelTypes]);

  // When activeChartId changes, refetch data
  useEffect(() => {
    if (activeChartId !== null) {
      chartRef.current?.refreshData();
      fetchAnnotations(activeChartId);
      fetchSpanAnnotations(activeChartId);
      setSelectedLabelId(null);
      setSelectedSpanId(null);
    }
  }, [activeChartId, fetchAnnotations, fetchSpanAnnotations]);

  const handleExport = () => {
    if (activeChartId) {
      window.location.href = `/api/export?chartId=${activeChartId}`;
    } else {
      window.location.href = '/api/export';
    }
  };

  const handleUploadSuccess = (chart: { id: number; name: string }) => {
    // Add new chart to list and select it
    const newChart: Chart = {
      id: chart.id,
      name: chart.name,
      created_at: Math.floor(Date.now() / 1000),
    };
    setCharts((prev) => [newChart, ...prev]);
    setActiveChartId(chart.id);
  };

  const handleAnnotationChange = async () => {
    await chartRef.current?.refreshData();
    await fetchAnnotations(activeChartId);
  };

  const handleSpanAnnotationsChange = async () => {
    await fetchSpanAnnotations(activeChartId);
  };

  const handleSelectedSpanChange = (spanId: number | null) => {
    setSelectedSpanId(spanId);
  };

  const handleDeleteSpan = async (spanId: number) => {
    try {
      const response = await fetch(`/api/span-annotations/${spanId}`, {
        method: 'DELETE',
      });
      if (response.ok) {
        await fetchSpanAnnotations(activeChartId);
        if (selectedSpanId === spanId) {
          setSelectedSpanId(null);
        }
      }
    } catch (error) {
      console.error('Failed to delete span annotation:', error);
    }
  };

  const handleLabelDelete = async (id: number) => {
    setAnnotations(annotations.filter((a) => a.id !== id));
    if (selectedLabelId === id) {
      setSelectedLabelId(null);
    }
  };

  const handleLabelSelect = (id: number) => {
    setSelectedLabelId(id === -1 ? null : id);
  };

  const handleSelectChart = (chartId: number) => {
    setActiveChartId(chartId);
  };

  const handleDeleteChart = async (chartId: number) => {
    try {
      const response = await fetch(`/api/charts/${chartId}`, { method: 'DELETE' });
      if (response.ok) {
        const remaining = charts.filter((c) => c.id !== chartId);
        setCharts(remaining);
        if (activeChartId === chartId) {
          setActiveChartId(remaining.length > 0 ? remaining[0].id : null);
        }
      }
    } catch (error) {
      console.error('Failed to delete chart:', error);
    }
  };

  // Fetch model info and initialize selected labels
  const fetchModelInfo = useCallback(async () => {
    try {
      const response = await fetch('/api/model/info');
      if (!response.ok) {
        setIsModelOnline(false);
        throw new Error('Model info unavailable');
      }
      const data: ModelInfoResponse = await response.json();
      setIsModelOnline(true);
      setPredictionState((prev) => ({
        ...prev,
        modelInfo: data,
        selectedLabels: new Set(data.label_config.map((l) => l.name)),
        error: null,
      }));
      return data;
    } catch (error) {
      console.error('Failed to fetch model info:', error);
      setIsModelOnline(false);
      setPredictionState((prev) => ({
        ...prev,
        modelInfo: null,
        error: error instanceof Error ? error.message : 'Failed to fetch model info',
      }));
      return null;
    }
  }, []);

  // Generate cache key from chart, timerange, and model version
  const generateCacheKey = useCallback((chartId: number | null, modelVersion?: string) => {
    if (!chartId) return null;
    const version = modelVersion || predictionState.modelInfo?.model_info.model_version || 'unknown';
    return `${chartId}_${version}`;
  }, [predictionState.modelInfo]);

  // Fetch predictions for visible candles
  const fetchPredictions = useCallback(async (candles: any[]) => {
    if (!activeChartId || candles.length === 0) return;

    const cacheKey = generateCacheKey(activeChartId, predictionState.modelInfo?.model_info.model_version);
    
    // Check cache first
    if (cacheKey && predictionCacheRef.current.has(cacheKey)) {
      const cached = predictionCacheRef.current.get(cacheKey)!;
      if (cached.modelVersion === predictionState.modelInfo?.model_info.model_version) {
        setPredictionState((prev) => ({
          ...prev,
          spans: cached.spans,
          perCandlePredictions: cached.predictions,
          cacheKey,
        }));
        return;
      }
    }

    setPredictionState((prev) => ({ ...prev, isLoading: true, error: null }));

    try {
      const response = await fetch('/api/predict', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ candles }),
      });

      if (!response.ok) {
        throw new Error(`Prediction failed: ${response.statusText}`);
      }

      const data = await response.json();
      
      // Cache the results
      if (cacheKey) {
        predictionCacheRef.current.set(cacheKey, {
          spans: data.spans,
          predictions: data.predictions,
          modelVersion: data.model_info.model_version,
        });
      }

      setPredictionState((prev) => ({
        ...prev,
        spans: data.spans,
        perCandlePredictions: data.predictions,
        isLoading: false,
        cacheKey,
      }));
    } catch (error) {
      console.error('Failed to fetch predictions:', error);
      setPredictionState((prev) => ({
        ...prev,
        isLoading: false,
        error: error instanceof Error ? error.message : 'Failed to fetch predictions',
      }));
    }
  }, [activeChartId, predictionState.modelInfo, generateCacheKey]);

  // Toggle prediction visibility
  const togglePredictionVisibility = useCallback(() => {
    setPredictionState((prev) => ({ ...prev, visible: !prev.visible }));
  }, []);

  // Update confidence threshold
  const setConfidenceThreshold = useCallback((threshold: number) => {
    setPredictionState((prev) => ({ ...prev, confidenceThreshold: threshold }));
  }, []);

  // Toggle label selection
  const toggleLabelSelection = useCallback((label: string) => {
    setPredictionState((prev) => {
      const newSelected = new Set(prev.selectedLabels);
      if (newSelected.has(label)) {
        newSelected.delete(label);
      } else {
        newSelected.add(label);
      }
      return { ...prev, selectedLabels: newSelected };
    });
  }, []);

  // Toggle show only disagreements
  const toggleShowOnlyDisagreements = useCallback(() => {
    setShowOnlyDisagreements((prev) => !prev);
  }, []);

  // Handle on-demand prediction for visible candles
  const handleFetchVisiblePredictions = useCallback(() => {
    // This will be called by the PredictionPanel
    // The actual candles data will be fetched from the chart ref
    const candles = chartRef.current?.getVisibleCandles();
    if (candles && candles.length > 0) {
      fetchPredictions(candles);
    }
  }, [fetchPredictions]);

  // Handle batch prediction for all candles
  const handleFetchBatchPredictions = useCallback(async () => {
    if (!activeChartId) return;

    setPredictionState((prev) => ({ ...prev, isLoading: true, error: null }));

    try {
      // Fetch chart data to get pair/timeframe info
      const chartResponse = await fetch(`/api/charts/${activeChartId}`);
      if (!chartResponse.ok) {
        throw new Error('Failed to fetch chart info');
      }
      const chartData = await chartResponse.json();

      // Fetch candles for the chart
      const candlesResponse = await fetch(`/api/candles?chartId=${activeChartId}`);
      if (!candlesResponse.ok) {
        throw new Error('Failed to fetch candles');
      }
      const candlesData = await candlesResponse.json();
      
      if (candlesData.length === 0) {
        throw new Error('No candles found for this chart');
      }

      const startTime = candlesData[0].time;
      const endTime = candlesData[candlesData.length - 1].time;

      // Make batch prediction request
      const response = await fetch('/api/predict/batch', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          pair: chartData.name,
          timeframe: '1h', // TODO: Get from chart metadata
          start_time: startTime,
          end_time: endTime,
        }),
      });

      if (!response.ok) {
        throw new Error(`Batch prediction failed: ${response.statusText}`);
      }

      const data = await response.json();
      
      const cacheKey = generateCacheKey(activeChartId, data.model_info.model_version);
      if (cacheKey) {
        predictionCacheRef.current.set(cacheKey, {
          spans: data.spans,
          predictions: data.predictions,
          modelVersion: data.model_info.model_version,
        });
      }

      setPredictionState((prev) => ({
        ...prev,
        spans: data.spans,
        perCandlePredictions: data.predictions,
        isLoading: false,
        cacheKey,
      }));
    } catch (error) {
      console.error('Failed to fetch batch predictions:', error);
      setPredictionState((prev) => ({
        ...prev,
        isLoading: false,
        error: error instanceof Error ? error.message : 'Failed to fetch batch predictions',
      }));
    }
  }, [activeChartId, generateCacheKey]);

  // Clear prediction cache when model version changes
  useEffect(() => {
    if (predictionState.modelInfo) {
      const currentVersion = predictionState.modelInfo.model_info.model_version;
      // Clear cache entries with different model versions
      const newCache = new Map();
      for (const [key, value] of predictionCacheRef.current.entries()) {
        if (value.modelVersion === currentVersion) {
          newCache.set(key, value);
        }
      }
      predictionCacheRef.current = newCache;
    }
  }, [predictionState.modelInfo?.model_info.model_version]);

  // Health polling - check model status every 30 seconds when offline
  useEffect(() => {
    if (!isModelOnline) {
      const interval = setInterval(() => {
        fetchModelInfo();
      }, 30000);
      return () => clearInterval(interval);
    }
  }, [isModelOnline, fetchModelInfo]);

  // Initialize model info on mount
  useEffect(() => {
    fetchModelInfo();
  }, [fetchModelInfo]);

  // Compute prediction summary when predictions or span annotations change
  useEffect(() => {
    if (predictionState.visible && predictionState.spans.length > 0) {
      const summary = detectDisagreements(spanAnnotations, predictionState.spans);
      setPredictionSummary(summary);
    } else {
      setPredictionSummary(null);
    }
  }, [predictionState.visible, predictionState.spans, spanAnnotations]);

  // Keyboard handler for Delete/Backspace key
  useEffect(() => {
    const handleKeyDown = async (e: KeyboardEvent) => {
      if ((e.key === 'Delete' || e.key === 'Backspace') && selectedLabelId !== null) {
        try {
          const response = await fetch(`/api/annotations/${selectedLabelId}`, {
            method: 'DELETE',
          });
          if (response.ok) {
            setSelectedLabelId(null);
            chartRef.current?.refreshData();
          }
        } catch (error) {
          console.error('Failed to delete label:', error);
        }
      }
    };

    window.addEventListener('keydown', handleKeyDown);
    return () => window.removeEventListener('keydown', handleKeyDown);
  }, [selectedLabelId]);

  return (
    <div className="flex h-screen bg-background">
      {/* Sidebar */}
      <aside className="w-72 flex-shrink-0 flex flex-col border-r border-border bg-card">
        <div className="p-6 border-b border-border">
          <h1 className="text-2xl font-semibold text-foreground">Candle Annotator</h1>
          <p className="text-sm text-muted-foreground mt-1">Chart annotation tool</p>
          <div className="mt-3 flex flex-col gap-1">
            <a
              href="/annotation-types"
              className="text-sm text-muted-foreground hover:text-foreground"
            >
              Manage Annotation Types
            </a>
            <a
              href="/span-label-types"
              className="text-sm text-muted-foreground hover:text-foreground"
            >
              Manage Span Label Types
            </a>
          </div>
        </div>
        <div className="p-6 pb-3">
          <ChartSelector
            charts={charts}
            activeChartId={activeChartId}
            onSelectChart={handleSelectChart}
            onDeleteChart={handleDeleteChart}
          />
        </div>
        <div className="px-6 pb-3">
          <FileUpload onUploadSuccess={handleUploadSuccess} />
        </div>
        <div className="px-6 pb-6">
          <Toolbox
            activeTool={activeTool}
            onToolChange={setActiveTool}
            onExport={handleExport}
            selectedColor={selectedColor}
            onColorChange={setSelectedColor}
            annotations={annotations}
            selectedLabelId={selectedLabelId}
            onLabelSelect={handleLabelSelect}
            onLabelDelete={handleLabelDelete}
            activeChartId={activeChartId}
            spanAnnotations={spanAnnotations}
            spanLabelTypes={spanLabelTypes}
            selectedSpanId={selectedSpanId}
            onSelectSpan={handleSelectedSpanChange}
            onDeleteSpan={handleDeleteSpan}
          />
        </div>
        <PredictionPanel
          predictionState={predictionState}
          onToggleVisibility={togglePredictionVisibility}
          onFetchPredictions={handleFetchVisiblePredictions}
          onFetchBatchPredictions={handleFetchBatchPredictions}
          onConfidenceChange={setConfidenceThreshold}
          onToggleLabelSelection={toggleLabelSelection}
          predictionSummary={predictionSummary}
          isModelOnline={isModelOnline}
          showOnlyDisagreements={showOnlyDisagreements}
          onToggleShowOnlyDisagreements={toggleShowOnlyDisagreements}
        />
      </aside>

      {/* Main chart area */}
      <main className="flex-1 relative bg-background">
        {/* Loading overlay for predictions */}
        {predictionState.isLoading && (
          <div className="absolute inset-0 bg-background/50 backdrop-blur-sm z-50 flex items-center justify-center">
            <div className="bg-card border border-border rounded-lg p-6 shadow-lg">
              <div className="flex items-center gap-3">
                <div className="animate-spin rounded-full h-5 w-5 border-2 border-primary border-t-transparent" />
                <span className="text-sm text-foreground">Loading predictions...</span>
              </div>
            </div>
          </div>
        )}
        <CandleChart
          ref={chartRef}
          activeTool={activeTool}
          onAnnotationChange={handleAnnotationChange}
          selectedColor={selectedColor}
          selectedLabelId={selectedLabelId}
          onLabelSelect={handleLabelSelect}
          activeChartId={activeChartId}
          spanAnnotations={spanAnnotations}
          spanLabelTypes={spanLabelTypes}
          selectedSpanId={selectedSpanId}
          onSpanAnnotationsChange={handleSpanAnnotationsChange}
          onSelectedSpanChange={handleSelectedSpanChange}
          predictionVisible={predictionState.visible}
          perCandlePredictions={predictionState.perCandlePredictions}
          predictionSpans={predictionState.spans}
          confidenceThreshold={predictionState.confidenceThreshold}
          selectedLabels={predictionState.selectedLabels}
          modelInfo={predictionState.modelInfo}
          predictionSummary={predictionSummary}
          showOnlyDisagreements={showOnlyDisagreements}
        />
      </main>
    </div>
  );
}