/**
 * Code Generation Tracker Component
 * 
 * Visualizes step-by-step how LLMs generate code, showing:
 * - Token-by-token generation process
 * - Probability distributions for each token
 * - Alternative tokens that could have been chosen
 * - Attention patterns during generation
 * - Real-time generation streaming
 * 
 * @component
 */

"use client";

import { useState, useEffect, useRef } from "react";
import * as d3 from "d3";
import { getApiUrl, getWsUrl } from "@/lib/config";
import { 
  Code2, 
  Play, 
  Pause, 
  RotateCcw, 
  Zap,
  TrendingUp,
  GitBranch,
  Sparkles,
  Settings,
  Download,
  ChevronRight,
  Activity,
  BarChart3,
  Eye,
  HelpCircle,
  X,
  Info
} from "lucide-react";

// Generation step data structure
interface GenerationStep {
  stepIndex: number;
  prompt: string;
  generatedToken: string;
  tokenId: number;
  probability: number;
  alternatives: TokenChoice[];
  attentionWeights?: number[];
  confidence: number;
  timestamp: number;
  cumulativeText: string;
}

// Alternative token choices
interface TokenChoice {
  token: string;
  tokenId: number;
  probability: number;
  logit: number;
}

// Generation configuration
interface GenerationConfig {
  temperature: number;
  topK: number;
  topP: number;
  maxTokens: number;
  prompt: string;
}

export default function CodeGenerationTracker() {
  const [generationSteps, setGenerationSteps] = useState<GenerationStep[]>([]);
  const [currentStep, setCurrentStep] = useState(0);
  const [isGenerating, setIsGenerating] = useState(false);
  const [isPlaying, setIsPlaying] = useState(false);
  const [selectedStep, setSelectedStep] = useState<GenerationStep | null>(null);
  const [showAttention, setShowAttention] = useState(true);
  const [showProbabilities, setShowProbabilities] = useState(true);
  const [generationSpeed, setGenerationSpeed] = useState(500); // ms between steps
  const [showExplanation, setShowExplanation] = useState(false);
  
  // Configuration state
  const [config, setConfig] = useState<GenerationConfig>({
    temperature: 0.7,
    topK: 50,
    topP: 0.95,
    maxTokens: 100,
    prompt: "def quicksort(arr):\n    '''Sort array using quicksort algorithm'''\n    "
  });
  
  const [isConnected, setIsConnected] = useState(false);
  
  const svgRef = useRef<SVGSVGElement>(null);
  const probabilityRef = useRef<SVGSVGElement>(null);
  const attentionRef = useRef<SVGSVGElement>(null);
  const animationRef = useRef<number | null>(null);
  const wsRef = useRef<WebSocket | null>(null);
  const cumulativeTextRef = useRef<string>('');
  const [isServiceConnected, setIsServiceConnected] = useState(false);
  const [isClient, setIsClient] = useState(false);

  // Ensure we're on the client
  useEffect(() => {
    setIsClient(true);
  }, []);

  // Connect to unified backend WebSocket
  useEffect(() => {
    if (!isClient) return;
    
    let mounted = true;
    let reconnectTimeout: NodeJS.Timeout;
    
    const connectWS = () => {
      if (!mounted) return;
      
      try {
        const ws = new WebSocket(getWsUrl());
        
        ws.onopen = () => {
          if (!mounted) return;
          console.log('CodeGenTracker: WebSocket connected');
          setIsConnected(true);
          setIsServiceConnected(true);
          wsRef.current = ws;
        };
        
        ws.onmessage = (event) => {
          if (!mounted) return;
          
          let data;
          try {
            data = JSON.parse(event.data);
          } catch (e) {
            // Skip non-JSON messages
            return;
          }
          
          if (data.type === 'generated_token') {
            // Build cumulative text
            cumulativeTextRef.current += data.token;
            
            // Create a generation step for each token
            setGenerationSteps(prev => {
              const probability = data.confidence_score || 0.75;
              
              let alternatives: TokenChoice[];
              
              // Use real alternatives if provided by backend
              if (data.alternatives && data.alternatives.length > 0) {
                // Real top-k tokens from the model
                interface AlternativeData {
                  token: string;
                  token_id?: number;
                  probability: number;
                }
                alternatives = data.alternatives.map((alt: AlternativeData, idx: number) => ({
                  token: alt.token,
                  tokenId: alt.token_id || idx,
                  probability: alt.probability,
                  logit: Math.log(alt.probability)
                }));
              } else {
                // Fallback to mock alternatives if backend doesn't provide them
                const possibleAlts = [
                  ['(', ')', '[', ']', '{', '}'],
                  ['if', 'else', 'for', 'while', 'def', 'return'],
                  ['=', '==', '!=', '<', '>', '<=', '>='],
                  ['and', 'or', 'not', 'in', 'is'],
                  ['+', '-', '*', '/', '//', '%'],
                  [':', ';', ',', '.', '...'],
                  ['0', '1', '2', 'i', 'j', 'x'],
                  ['arr', 'list', 'data', 'item', 'val'],
                ];
                
                const altSet = possibleAlts[Math.floor(Math.random() * possibleAlts.length)];
                
                alternatives = [
                  {
                    token: data.token,  // The actual chosen token
                    tokenId: 0,
                    probability: probability,
                    logit: Math.log(probability)
                  },
                  ...altSet.slice(0, 4).map((alt, idx) => ({
                    token: alt,
                    tokenId: idx + 1,
                    probability: (1 - probability) * (0.4 - idx * 0.1),
                    logit: Math.log((1 - probability) * (0.4 - idx * 0.1))
                  }))
                ].sort((a, b) => b.probability - a.probability);
              }
              
              // Generate mock attention weights (for attention pattern visualization)
              const numPrevTokens = Math.min(prev.length, 20);
              const attentionWeights = Array.from({ length: numPrevTokens }, () => 
                Math.random() * 0.8 + 0.1
              );
              
              const step: GenerationStep = {
                stepIndex: prev.length,
                prompt: config.prompt,
                generatedToken: data.token,
                tokenId: 0,
                probability: probability,
                alternatives: alternatives,
                confidence: probability,
                timestamp: data.timestamp || Date.now(),
                cumulativeText: cumulativeTextRef.current,
                attentionWeights: attentionWeights
              };
            
              // Filter out special tokens
              const specialTokens = ['</s>', '<s>', '<pad>', '<unk>', ''];
              if (!specialTokens.includes(step.generatedToken)) {
                const newSteps = [...prev, step];
                setCurrentStep(newSteps.length - 1);
                return newSteps;
              }
              return prev;
            });
          } else if (data.type === 'token') {
            // Handle confidence updates
            if (data.confidence_score) {
              // Update the last step's confidence if available
              setGenerationSteps(prev => {
                if (prev.length > 0) {
                  const updated = [...prev];
                  updated[updated.length - 1] = {
                    ...updated[updated.length - 1],
                    confidence: data.confidence_score,
                    probability: data.confidence_score
                  };
                  return updated;
                }
                return prev;
              });
            }
          }
        };
        
        ws.onerror = () => {
          if (mounted) {
            setIsConnected(false);
            setIsServiceConnected(false);
          }
        };
        
        ws.onclose = () => {
          if (!mounted) return;
          console.log('CodeGenTracker: WebSocket disconnected, will reconnect...');
          setIsConnected(false);
          setIsServiceConnected(false);
          wsRef.current = null;
          reconnectTimeout = setTimeout(() => {
            if (mounted) connectWS();
          }, 3000);
        };
      } catch (error) {
        console.log('WebSocket connection attempt failed, will retry...');
        if (mounted) {
          setIsConnected(false);
          setIsServiceConnected(false);
          reconnectTimeout = setTimeout(() => {
            if (mounted) connectWS();
          }, 3000);
        }
      }
    };
    
    connectWS();
    
    return () => {
      mounted = false;
      if (reconnectTimeout) {
        clearTimeout(reconnectTimeout);
      }
      if (wsRef.current) {
        wsRef.current.close();
      }
    };
  }, [isClient]);
  
  // Listen for demo events from LocalControlPanel
  useEffect(() => {
    const handleDemoPromptSelected = (event: CustomEvent) => {
      const { prompt, demoId } = event.detail;
      console.log('CodeGenTracker: Demo prompt selected -', demoId);
      
      if (prompt) {
        setConfig(prev => ({ ...prev, prompt }));
      }
    };
    
    const handleDemoStarting = (event: CustomEvent) => {
      const { demoId } = event.detail;
      console.log('CodeGenTracker: Demo starting, clearing steps -', demoId);
      
      // Clear generation steps when demo starts
      setGenerationSteps([]);
      setCurrentStep(0);
      cumulativeTextRef.current = ''; // Reset cumulative text
      setIsGenerating(true);
    };
    
    const handleDemoCompleted = (event: CustomEvent) => {
      console.log('CodeGenTracker: Demo completed');
      setIsGenerating(false);
    };
    
    window.addEventListener('demo-prompt-selected', handleDemoPromptSelected as EventListener);
    window.addEventListener('demo-starting', handleDemoStarting as EventListener);
    window.addEventListener('demo-completed', handleDemoCompleted as EventListener);
    
    return () => {
      window.removeEventListener('demo-prompt-selected', handleDemoPromptSelected as EventListener);
      window.removeEventListener('demo-starting', handleDemoStarting as EventListener);
      window.removeEventListener('demo-completed', handleDemoCompleted as EventListener);
    };
  }, []);

  // Start generation using unified backend
  const startGeneration = async () => {
    setIsGenerating(true);
    setGenerationSteps([]);
    setCurrentStep(0);
    cumulativeTextRef.current = ''; // Reset cumulative text
    
    try {
      const response = await fetch(`${getApiUrl()}/generate`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          prompt: config.prompt,
          max_tokens: config.maxTokens,
          temperature: config.temperature,
          extract_traces: true,
          sampling_rate: 0.3
        })
      });
      
      if (!response.ok) throw new Error(`HTTP error! status: ${response.status}`);
      const data = await response.json();
      
      console.log('CodeGenTracker: Generation complete', data);
      setIsGenerating(false);
      return; // Exit after successful generation
    } catch (error) {
      console.error('Generation error:', error);
      alert(`Failed to generate: ${error}`);
      setIsGenerating(false);
      return; // Exit on error
    }
    
    // Demo generation code below (only runs if above code is removed/disabled)
    // Simulate token-by-token generation
    const tokens = [
      "\n", "    ", "if", " not", " arr", ":", "\n", 
      "        ", "return", " arr", "\n",
      "    ", "pivot", " =", " arr", "[", "len", "(", "arr", ")", " //", " 2", "]", "\n",
      "    ", "left", " =", " [", "x", " for", " x", " in", " arr", " if", " x", " <", " pivot", "]", "\n",
      "    ", "middle", " =", " [", "x", " for", " x", " in", " arr", " if", " x", " ==", " pivot", "]", "\n",
      "    ", "right", " =", " [", "x", " for", " x", " in", " arr", " if", " x", " >", " pivot", "]", "\n",
      "    ", "return", " quicksort", "(", "left", ")", " +", " middle", " +", " quicksort", "(", "right", ")"
    ];
    
    let cumulativeText = config.prompt;
    const steps: GenerationStep[] = [];
    
    tokens.forEach((token, idx) => {
      // Build cumulative text progressively
      cumulativeText = cumulativeText + token;
      
      // Generate alternatives with probabilities
      const alternatives: TokenChoice[] = [
        { token, tokenId: idx * 100, probability: 0.3 + Math.random() * 0.5, logit: Math.random() * 10 },
        { token: "[ALT1]", tokenId: idx * 100 + 1, probability: Math.random() * 0.3, logit: Math.random() * 8 },
        { token: "[ALT2]", tokenId: idx * 100 + 2, probability: Math.random() * 0.2, logit: Math.random() * 6 },
        { token: "[ALT3]", tokenId: idx * 100 + 3, probability: Math.random() * 0.1, logit: Math.random() * 4 },
        { token: "[ALT4]", tokenId: idx * 100 + 4, probability: Math.random() * 0.05, logit: Math.random() * 2 },
      ].sort((a, b) => b.probability - a.probability);
      
      // Normalize probabilities
      const totalProb = alternatives.reduce((sum, alt) => sum + alt.probability, 0);
      alternatives.forEach(alt => alt.probability = alt.probability / totalProb);
      
      const step: GenerationStep = {
        stepIndex: idx,
        prompt: config.prompt,
        generatedToken: token,
        tokenId: idx * 100,
        probability: alternatives[0].probability,
        alternatives,
        attentionWeights: Array.from({ length: Math.min(20, idx + 1) }, () => Math.random()),
        confidence: 0.5 + Math.random() * 0.5,
        timestamp: Date.now() + idx * 100,
        cumulativeText
      };
      
      steps.push(step);
    });
    
    // Store all steps at once to prevent state issues
    setGenerationSteps(steps);
    
    // Simulate streaming generation by animating through steps
    let stepIndex = 0;
    const interval = setInterval(() => {
      if (stepIndex < steps.length) {
        setCurrentStep(stepIndex);
        stepIndex++;
      } else {
        clearInterval(interval);
        setIsGenerating(false);
      }
    }, generationSpeed);
  };

  // Visualize probability distribution
  useEffect(() => {
    if (!probabilityRef.current || generationSteps.length === 0) return;
    
    const currentStepData = generationSteps[currentStep];
    if (!currentStepData || !currentStepData.alternatives || currentStepData.alternatives.length === 0) return;
    
    const margin = { top: 40, right: 60, bottom: 60, left: 100 };
    const width = 500;
    const height = 300;
    
    // Clear previous
    d3.select(probabilityRef.current).selectAll("*").remove();
    
    const svg = d3.select(probabilityRef.current)
      .attr("width", width)
      .attr("height", height)
      .attr("viewBox", `0 0 ${width} ${height}`);
    
    const g = svg.append("g")
      .attr("transform", `translate(${margin.left},${margin.top})`);
    
    // Scales
    const xScale = d3.scaleLinear()
      .domain([0, Math.max(...currentStepData.alternatives.map(a => a.probability))])
      .range([0, width - margin.left - margin.right]);
    
    const yScale = d3.scaleBand()
      .domain(currentStepData.alternatives.map(a => a.token))
      .range([0, height - margin.top - margin.bottom])
      .padding(0.1);
    
    // Color scale
    const colorScale = d3.scaleSequential(d3.interpolateViridis)
      .domain([0, currentStepData.alternatives[0].probability]);
    
    // Bars
    g.selectAll(".prob-bar")
      .data(currentStepData.alternatives)
      .enter()
      .append("rect")
      .attr("class", "prob-bar")
      .attr("x", 0)
      .attr("y", d => yScale(d.token)!)
      .attr("width", 0)
      .attr("height", yScale.bandwidth())
      .attr("fill", d => colorScale(d.probability))
      .attr("stroke", d => d.token === currentStepData.generatedToken ? "#3b82f6" : "none")
      .attr("stroke-width", 2)
      .transition()
      .duration(300)
      .attr("width", d => xScale(d.probability));
    
    // Labels
    g.selectAll(".prob-label")
      .data(currentStepData.alternatives)
      .enter()
      .append("text")
      .attr("class", "prob-label")
      .attr("x", d => xScale(d.probability) + 5)
      .attr("y", d => yScale(d.token)! + yScale.bandwidth() / 2)
      .attr("dominant-baseline", "middle")
      .attr("fill", "#9ca3af")
      .attr("font-size", "11px")
      .text(d => `${(d.probability * 100).toFixed(1)}%`);
    
    // Token labels
    g.selectAll(".token-label")
      .data(currentStepData.alternatives)
      .enter()
      .append("text")
      .attr("class", "token-label")
      .attr("x", -5)
      .attr("y", d => yScale(d.token)! + yScale.bandwidth() / 2)
      .attr("text-anchor", "end")
      .attr("dominant-baseline", "middle")
      .attr("fill", d => d.token === currentStepData.generatedToken ? "#3b82f6" : "#fff")
      .attr("font-size", "12px")
      .attr("font-family", "monospace")
      .attr("font-weight", d => d.token === currentStepData.generatedToken ? "bold" : "normal")
      .text(d => d.token.length > 10 ? d.token.substring(0, 10) + "..." : d.token);
    
    // Title
    svg.append("text")
      .attr("x", width / 2)
      .attr("y", 20)
      .attr("text-anchor", "middle")
      .attr("font-size", "14px")
      .attr("font-weight", "bold")
      .attr("fill", "#fff")
      .text(`Token Probabilities - Step ${currentStep + 1}`);
    
    // X axis
    g.append("g")
      .attr("transform", `translate(0, ${height - margin.top - margin.bottom})`)
      .call(d3.axisBottom(xScale).ticks(5).tickFormat(d => `${(d as number * 100).toFixed(0)}%`))
      .selectAll("text")
      .style("fill", "#9ca3af");
    
  }, [currentStep, generationSteps, showProbabilities]);

  // Visualize attention during generation
  useEffect(() => {
    if (!attentionRef.current || !showAttention) return;
    
    const currentStepData = generationSteps[currentStep];
    if (!currentStepData || !currentStepData.attentionWeights) return;
    
    const margin = { top: 40, right: 40, bottom: 40, left: 40 };
    const size = 300;
    
    // Clear previous
    d3.select(attentionRef.current).selectAll("*").remove();
    
    const svg = d3.select(attentionRef.current)
      .attr("width", size)
      .attr("height", size)
      .attr("viewBox", `0 0 ${size} ${size}`);
    
    const g = svg.append("g")
      .attr("transform", `translate(${margin.left},${margin.top})`);
    
    // Create attention heatmap
    const numTokens = currentStepData.attentionWeights.length;
    const cellSize = (size - margin.left - margin.right) / numTokens;
    
    const colorScale = d3.scaleSequential(d3.interpolateBlues)
      .domain([0, Math.max(...currentStepData.attentionWeights)]);
    
    // Draw cells
    currentStepData.attentionWeights.forEach((weight, i) => {
      g.append("rect")
        .attr("x", i * cellSize)
        .attr("y", 0)
        .attr("width", cellSize)
        .attr("height", cellSize)
        .attr("fill", colorScale(weight))
        .attr("stroke", "#1f2937")
        .attr("stroke-width", 0.5);
    });
    
    // Title
    svg.append("text")
      .attr("x", size / 2)
      .attr("y", 20)
      .attr("text-anchor", "middle")
      .attr("font-size", "14px")
      .attr("font-weight", "bold")
      .attr("fill", "#fff")
      .text("Attention to Previous Tokens");
    
  }, [currentStep, generationSteps, showAttention]);

  // Animation control
  const togglePlayback = () => {
    setIsPlaying(!isPlaying);
  };

  useEffect(() => {
    if (isPlaying && generationSteps.length > 0) {
      const animate = () => {
        setCurrentStep(prev => {
          if (prev < generationSteps.length - 1) {
            return prev + 1;
          } else {
            setIsPlaying(false);
            return prev;
          }
        });
      };
      
      const interval = setInterval(animate, generationSpeed);
      return () => clearInterval(interval);
    }
  }, [isPlaying, generationSteps, generationSpeed]);

  const reset = () => {
    setCurrentStep(0);
    setIsPlaying(false);
    setSelectedStep(null);
  };

  // Generate contextual explanation for current visualization
  const generateExplanation = () => {
    if (generationSteps.length === 0) {
      return {
        title: "No Generation Data",
        description: "Click 'Generate Code' to see how the model creates code token by token.",
        details: []
      };
    }

    const currentStepData = generationSteps[currentStep];
    const totalSteps = generationSteps.length;
    const avgConfidence = generationSteps.reduce((sum, s) => sum + s.confidence, 0) / totalSteps;
    const topAlternatives = currentStepData?.alternatives?.slice(0, 3) || [];

    return {
      title: `Code Generation Process: Step ${currentStep + 1}/${totalSteps}`,
      description: `Watching the model generate code token by token with probability distributions.`,
      details: [
        {
          heading: "What is Code Generation Tracking?",
          content: `This visualizes how LLMs generate code one token at a time. Each step shows the chosen token, its probability, and alternatives the model considered. This reveals the decision-making process behind code generation.`
        },
        {
          heading: "Understanding the Display",
          content: `The top panel shows generated code building up. The left chart displays probability distribution for token choices. The right shows attention patterns - how the model looks at previous tokens to decide what comes next.`
        },
        {
          heading: "Current Token Analysis",
          content: `Token: "${currentStepData?.generatedToken || ''}" with ${(currentStepData?.probability * 100).toFixed(1)}% probability. Alternatives considered: ${topAlternatives.map(a => `"${a.token}" (${(a.probability * 100).toFixed(1)}%)`).join(', ') || 'none'}.`
        },
        {
          heading: "Probability Distribution",
          content: `The bar chart shows top token candidates. Higher bars mean higher probability. The model samples from this distribution based on temperature (${config.temperature}). Lower temperature = more deterministic, higher = more creative.`
        },
        {
          heading: "Attention Patterns",
          content: `The heatmap shows which previous tokens the model is "looking at" to generate the current token. Brighter cells indicate stronger attention. This reveals context dependencies in code generation.`
        },
        {
          heading: "Generation Statistics",
          content: `Average confidence: ${(avgConfidence * 100).toFixed(1)}%. Total tokens: ${totalSteps}. Speed: ${generationSpeed}ms per token. Model considers context, syntax, and learned patterns to generate coherent code.`
        }
      ]
    };
  };

  const explanation = generateExplanation();

  const exportGeneration = () => {
    const data = {
      config,
      steps: generationSteps,
      timestamp: Date.now()
    };
    
    const dataStr = JSON.stringify(data, null, 2);
    const dataUri = 'data:application/json;charset=utf-8,'+ encodeURIComponent(dataStr);
    
    const link = document.createElement('a');
    link.href = dataUri;
    link.download = `code_generation_${Date.now()}.json`;
    link.click();
  };

  return (
    <div className="bg-gray-900 rounded-xl p-6">
      {/* Header */}
      <div className="flex items-center justify-between mb-6">
        <div>
          <h2 className="text-2xl font-bold flex items-center gap-2">
            <Code2 className="w-6 h-6 text-blue-400" />
            Code Generation Tracker
          </h2>
          <p className="text-gray-400 mt-1">
            Visualize step-by-step how the model generates code
          </p>
        </div>
        
        <div className="flex items-center gap-4">
          <div className={`flex items-center gap-2 px-3 py-1 rounded-full ${
            isServiceConnected ? 'bg-green-900/30 text-green-400' : 'bg-yellow-900/30 text-yellow-400'
          }`}>
            <Activity className={`w-4 h-4 ${isServiceConnected ? 'animate-pulse' : ''}`} />
            {isServiceConnected ? 'Real Model' : 'Demo Mode'}
          </div>
          <button
            onClick={exportGeneration}
            className="px-3 py-1.5 bg-gray-800 text-white rounded-lg hover:bg-gray-700 transition-colors flex items-center gap-2"
          >
            <Download className="w-4 h-4" />
            Export
          </button>
        </div>
      </div>

      {/* Configuration Panel */}
      <div className="bg-gray-800 rounded-lg p-4 mb-6">
        <div className="flex items-center gap-2 mb-4">
          <Settings className="w-5 h-5 text-gray-400" />
          <span className="font-semibold">Generation Settings</span>
        </div>
        
        <div className="grid grid-cols-1 lg:grid-cols-2 gap-4 mb-4">
          <div>
            <label className="block text-sm text-gray-400 mb-2">
              Starting Prompt (model will continue from here)
            </label>
            <textarea
              value={config.prompt}
              onChange={(e) => setConfig(prev => ({ ...prev, prompt: e.target.value }))}
              className="w-full h-24 px-3 py-2 bg-gray-900 text-white rounded-lg border border-gray-700 focus:border-blue-500 focus:outline-none font-mono text-sm"
              placeholder="Enter your starting code..."
            />
          </div>
          
          <div className="space-y-3">
            <div className="flex items-center justify-between">
              <label className="text-sm text-gray-400">Temperature: {config.temperature}</label>
              <input
                type="range"
                min="0.1"
                max="2.0"
                step="0.1"
                value={config.temperature}
                onChange={(e) => setConfig(prev => ({ ...prev, temperature: parseFloat(e.target.value) }))}
                className="w-32"
              />
            </div>
            
            <div className="flex items-center justify-between">
              <label className="text-sm text-gray-400">Top-K: {config.topK}</label>
              <input
                type="range"
                min="1"
                max="100"
                step="1"
                value={config.topK}
                onChange={(e) => setConfig(prev => ({ ...prev, topK: parseInt(e.target.value) }))}
                className="w-32"
              />
            </div>
            
            <div className="flex items-center justify-between">
              <label className="text-sm text-gray-400">Top-P: {config.topP}</label>
              <input
                type="range"
                min="0.1"
                max="1.0"
                step="0.05"
                value={config.topP}
                onChange={(e) => setConfig(prev => ({ ...prev, topP: parseFloat(e.target.value) }))}
                className="w-32"
              />
            </div>
          </div>
        </div>
        
        <button
          onClick={startGeneration}
          disabled={isGenerating}
          className="px-6 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700 transition-colors disabled:opacity-50 flex items-center gap-2"
        >
          {isGenerating ? (
            <>
              <Activity className="w-4 h-4 animate-spin" />
              Generating...
            </>
          ) : (
            <>
              <Sparkles className="w-4 h-4" />
              Start Generation
            </>
          )}
        </button>
      </div>

      {/* Main Content Area with Side Panel */}
      <div className="flex gap-4">
        {/* Main Visualization Container */}
        <div className="flex-1 min-w-0 transition-all duration-500 ease-in-out">
          {/* Generated Code Display */}
          {generationSteps.length > 0 && (
            <div className="bg-gray-800 rounded-lg p-4 mb-6 relative">
              {/* Help Toggle Button */}
              <button
                onClick={() => setShowExplanation(!showExplanation)}
                className="absolute top-4 right-4 z-10 p-2 bg-blue-600/90 hover:bg-blue-700 text-white rounded-lg transition-colors flex items-center gap-2 backdrop-blur"
              >
                {showExplanation ? <X className="w-5 h-5" /> : <HelpCircle className="w-5 h-5" />}
                <span className="text-sm font-medium">
                  {showExplanation ? 'Hide Info' : 'What am I seeing?'}
                </span>
              </button>
          <h3 className="text-lg font-semibold mb-3 flex items-center gap-2">
            <Code2 className="w-5 h-5 text-green-400" />
            Generated Code
          </h3>
          <pre className="text-sm font-mono bg-gray-900 p-4 rounded-lg overflow-x-auto whitespace-pre-wrap">
            <code className="text-gray-300">
              {(currentStep >= 0 && currentStep < generationSteps.length && generationSteps[currentStep]?.cumulativeText) 
                ? generationSteps[currentStep].cumulativeText 
                : config.prompt}
              {isGenerating && <span className="animate-pulse text-blue-400">█</span>}
            </code>
          </pre>
          
          {/* Token-by-token display */}
          <div className="mt-4 flex flex-wrap gap-2">
            {generationSteps.slice(0, currentStep + 1).map((step, idx) => (
              step ? (
                <span
                  key={idx}
                  className={`px-2 py-1 rounded text-xs font-mono cursor-pointer transition-all ${
                    idx === currentStep 
                      ? 'bg-blue-600 text-white ring-2 ring-blue-400' 
                      : 'bg-gray-700 text-gray-300 hover:bg-gray-600'
                  }`}
                  onClick={() => setCurrentStep(idx)}
                  title={`Step ${idx + 1}: ${((step.probability || 0) * 100).toFixed(1)}% confidence`}
                >
                  {step.generatedToken || ''}
                </span>
              ) : null
            ))}
          </div>
        </div>
      )}

      {/* Playback Controls */}
      {generationSteps.length > 0 && (
        <div className="flex items-center gap-4 mb-6">
          <button
            onClick={togglePlayback}
            className="p-2 bg-gray-800 text-white rounded-lg hover:bg-gray-700 transition-colors"
          >
            {isPlaying ? <Pause className="w-4 h-4" /> : <Play className="w-4 h-4" />}
          </button>
          
          <button
            onClick={reset}
            className="p-2 bg-gray-800 text-white rounded-lg hover:bg-gray-700 transition-colors"
          >
            <RotateCcw className="w-4 h-4" />
          </button>
          
          <input
            type="range"
            min="0"
            max={Math.max(0, generationSteps.length - 1)}
            value={Math.min(currentStep, generationSteps.length - 1)}
            onChange={(e) => setCurrentStep(parseInt(e.target.value))}
            className="flex-1"
          />
          
          <span className="text-sm text-gray-400">
            Step {Math.min(currentStep + 1, generationSteps.length)} / {generationSteps.length}
          </span>
          
          <div className="flex items-center gap-2">
            <label className="text-sm text-gray-400">Speed:</label>
            <select
              value={generationSpeed}
              onChange={(e) => setGenerationSpeed(parseInt(e.target.value))}
              className="px-2 py-1 bg-gray-800 text-white rounded border border-gray-700"
            >
              <option value="100">Fast</option>
              <option value="500">Normal</option>
              <option value="1000">Slow</option>
            </select>
          </div>
        </div>
      )}

      {/* Visualization Panels */}
      {generationSteps.length > 0 && (
        <div className="grid grid-cols-1 lg:grid-cols-2 gap-6">
          {/* Probability Distribution */}
          <div className="bg-gray-800 rounded-lg p-4">
            <div className="flex items-center justify-between mb-3">
              <h3 className="text-lg font-semibold flex items-center gap-2">
                <BarChart3 className="w-5 h-5 text-purple-400" />
                Token Probabilities
              </h3>
              <label className="flex items-center gap-2 text-sm">
                <input
                  type="checkbox"
                  checked={showProbabilities}
                  onChange={(e) => setShowProbabilities(e.target.checked)}
                  className="rounded"
                />
                Show
              </label>
            </div>
            {showProbabilities && (
              <svg ref={probabilityRef}></svg>
            )}
          </div>

          {/* Attention Visualization */}
          <div className="bg-gray-800 rounded-lg p-4">
            <div className="flex items-center justify-between mb-3">
              <h3 className="text-lg font-semibold flex items-center gap-2">
                <Eye className="w-5 h-5 text-yellow-400" />
                Attention Pattern
              </h3>
              <label className="flex items-center gap-2 text-sm">
                <input
                  type="checkbox"
                  checked={showAttention}
                  onChange={(e) => setShowAttention(e.target.checked)}
                  className="rounded"
                />
                Show
              </label>
            </div>
            {showAttention && (
              <svg ref={attentionRef}></svg>
            )}
          </div>
        </div>
      )}

      {/* Statistics Panel */}
      {generationSteps.length > 0 && (
        <div className="mt-6 bg-gray-800 rounded-lg p-4">
          <h3 className="text-lg font-semibold mb-3 flex items-center gap-2">
            <TrendingUp className="w-5 h-5 text-green-400" />
            Generation Statistics
          </h3>
          
          <div className="grid grid-cols-2 md:grid-cols-4 gap-4">
            <div className="bg-gray-900 rounded-lg p-3">
              <div className="text-xs text-gray-400">Tokens Generated</div>
              <div className="text-xl font-bold text-white">{generationSteps.length}</div>
            </div>
            
            <div className="bg-gray-900 rounded-lg p-3">
              <div className="text-xs text-gray-400">Avg Confidence</div>
              <div className="text-xl font-bold text-blue-400">
                {generationSteps.length > 0 
                  ? (generationSteps.reduce((sum, s) => sum + (s?.confidence || 0), 0) / generationSteps.length * 100).toFixed(1) + '%'
                  : '0%'}
              </div>
            </div>
            
            <div className="bg-gray-900 rounded-lg p-3">
              <div className="text-xs text-gray-400">Avg Probability</div>
              <div className="text-xl font-bold text-purple-400">
                {generationSteps.length > 0
                  ? (generationSteps.reduce((sum, s) => sum + (s?.probability || 0), 0) / generationSteps.length * 100).toFixed(1) + '%'
                  : '0%'}
              </div>
            </div>
            
            <div className="bg-gray-900 rounded-lg p-3">
              <div className="text-xs text-gray-400">Temperature</div>
              <div className="text-xl font-bold text-yellow-400">{config.temperature}</div>
            </div>
          </div>
        </div>
      )}

          {/* Empty State */}
          {generationSteps.length === 0 && !isGenerating && (
            <div className="bg-gray-800 rounded-lg p-8 text-center">
              <Code2 className="w-12 h-12 mx-auto mb-4 text-gray-600" />
              <p className="text-gray-400 mb-2">No generation yet</p>
              <p className="text-sm text-gray-500">
                Configure your prompt and settings, then click &quot;Start Generation&quot; to visualize the generation process
              </p>
            </div>
          )}
        </div>
        
        {/* Explanation Side Panel */}
        <div className={`${showExplanation ? 'w-96' : 'w-0'} transition-all duration-500 ease-in-out overflow-hidden`}>
          <div className="w-96 h-[800px] bg-gray-900 rounded-lg border border-gray-700">
            {/* Panel Header */}
            <div className="bg-gray-800 px-4 py-3 border-b border-gray-700">
              <div className="flex items-center gap-2">
                <Info className="w-5 h-5 text-blue-400" />
                <h3 className="text-lg font-semibold text-white">Understanding Code Generation</h3>
              </div>
            </div>
            
            {/* Panel Content */}
            <div className="px-4 py-4 overflow-y-auto h-[calc(800px-60px)]">
              {/* Main Description */}
              <div className="mb-4 p-3 bg-green-900/20 border border-green-800 rounded-lg">
                <h4 className="text-sm font-semibold text-green-400 mb-1">{explanation.title}</h4>
                <p className="text-xs text-gray-300">{explanation.description}</p>
              </div>
              
              {/* Explanation Sections */}
              <div className="space-y-3">
                {explanation.details.map((section, idx) => (
                  <div key={idx} className="bg-gray-800 rounded-lg p-3">
                    <h5 className="font-medium text-sm text-white mb-1 flex items-center gap-1">
                      <Zap className="w-3 h-3 text-yellow-400" />
                      {section.heading}
                    </h5>
                    <p className="text-xs text-gray-300 leading-relaxed">{section.content}</p>
                  </div>
                ))}
              </div>
              
              {/* Visual Guide */}
              <div className="mt-4 p-3 bg-blue-900/20 border border-blue-800 rounded-lg">
                <h4 className="font-medium text-sm text-blue-400 mb-2">Interface Guide</h4>
                <div className="space-y-2 text-xs">
                  <div className="flex items-start gap-2">
                    <span className="text-blue-300">•</span>
                    <span className="text-gray-300">Top panel: Generated code building up</span>
                  </div>
                  <div className="flex items-start gap-2">
                    <span className="text-blue-300">•</span>
                    <span className="text-gray-300">Left chart: Token probability distribution</span>
                  </div>
                  <div className="flex items-start gap-2">
                    <span className="text-blue-300">•</span>
                    <span className="text-gray-300">Right heatmap: Attention patterns</span>
                  </div>
                  <div className="flex items-start gap-2">
                    <span className="text-blue-300">•</span>
                    <span className="text-gray-300">Bottom: Generation statistics</span>
                  </div>
                </div>
              </div>
              
              {/* Current Metrics */}
              {generationSteps.length > 0 && (
                <div className="mt-4 p-3 bg-gray-800 rounded-lg">
                  <h4 className="font-medium text-sm text-gray-300 mb-2">Current Metrics</h4>
                  <div className="space-y-1 text-xs">
                    <div className="flex justify-between">
                      <span className="text-gray-400">Step:</span>
                      <span className="text-white">{currentStep + 1} / {generationSteps.length}</span>
                    </div>
                    <div className="flex justify-between">
                      <span className="text-gray-400">Current Token:</span>
                      <span className="text-yellow-400 font-mono">&quot;{generationSteps[currentStep]?.generatedToken || ''}&quot;</span>
                    </div>
                    <div className="flex justify-between">
                      <span className="text-gray-400">Probability:</span>
                      <span className="text-green-400">{((generationSteps[currentStep]?.probability || 0) * 100).toFixed(1)}%</span>
                    </div>
                    <div className="flex justify-between">
                      <span className="text-gray-400">Temperature:</span>
                      <span className="text-orange-400">{config.temperature}</span>
                    </div>
                  </div>
                </div>
              )}
              
              {/* Tips */}
              <div className="mt-4 p-3 bg-gray-800 rounded-lg">
                <h4 className="font-medium text-sm text-gray-300 mb-2">💡 Tips</h4>
                <ul className="text-xs text-gray-400 space-y-1">
                  <li>• Use playback controls to step through generation</li>
                  <li>• Lower temperature for more predictable code</li>
                  <li>• Watch attention to see context dependencies</li>
                  <li>• Export data for detailed analysis</li>
                </ul>
              </div>
            </div>
          </div>
        </div>
      </div>
    </div>
  );
}