frontend/ModelInspector.tsx

/**
 * Model Inspector Component
 * 
 * Displays detailed architecture information about the loaded model,
 * including layers, parameters, attention heads, and accessible components.
 * Makes the "black box" transparent by showing what can be visualized.
 * 
 * @component
 */

"use client";

import { useState, useEffect, lazy, Suspense } from "react";
import { getApiUrl } from "@/lib/config";
import { 
  Brain, 
  Layers, 
  Eye, 
  Cpu, 
  Database,
  GitBranch,
  Zap,
  ChevronRight,
  ChevronDown,
  Activity,
  Info,
  Box,
  Network,
  Move3D
} from "lucide-react";

// Lazy load the 3D components to avoid SSR issues
const ModelArchitecture3D = lazy(() => import('./ModelArchitecture3D'));
const DecisionPath3D = lazy(() => import('./DecisionPath3DEnhanced'));

interface ModelInfo {
  name: string;
  type: string;
  totalParams: number;
  layers: number;
  heads: number;
  hiddenSize: number;
  vocabSize: number;
  maxPositions: number;
  architecture: string[];
  accessible: string[];
}

export default function ModelInspector({ hideControlPanel = false }: { hideControlPanel?: boolean }) {
  const [expandedSections, setExpandedSections] = useState<Set<string>>(new Set(['overview']));
  const [isConnected, setIsConnected] = useState(false);
  const [isLoading, setIsLoading] = useState(true);
  const [modelInfo, setModelInfo] = useState<ModelInfo>({
    name: "Loading...",
    type: "unknown",
    totalParams: 0,
    layers: 0,
    heads: 0,
    hiddenSize: 0,
    vocabSize: 0,
    maxPositions: 0,
    architecture: [],
    accessible: []
  });
  const [deviceInfo, setDeviceInfo] = useState<string>("");
  const [modelConfig, setModelConfig] = useState<Record<string, unknown> | null>(null);

  // Fetch model information from backend
  useEffect(() => {
    const fetchModelInfo = async () => {
      try {
        const response = await fetch(`${getApiUrl()}/model/info`);
        if (response.ok) {
          const data = await response.json();
          setModelInfo({
            name: data.name,
            type: data.type,
            totalParams: data.totalParams,
            layers: data.layers,
            heads: data.heads,
            hiddenSize: data.hiddenSize,
            vocabSize: data.vocabSize,
            maxPositions: data.maxPositions,
            architecture: [data.architecture],
            accessible: data.accessible
          });
          setDeviceInfo(data.device);
          setModelConfig(data.config);
          setIsConnected(true);
        }
      } catch (error) {
        console.error('Failed to fetch model info:', error);
        // Keep default/mock data if fetch fails
        setIsConnected(false);
      } finally {
        setIsLoading(false);
      }
    };

    fetchModelInfo();
    // Refresh model info every 10 seconds
    const interval = setInterval(fetchModelInfo, 10000);
    return () => clearInterval(interval);
  }, []);

  const toggleSection = (section: string) => {
    const newExpanded = new Set(expandedSections);
    if (newExpanded.has(section)) {
      newExpanded.delete(section);
    } else {
      newExpanded.add(section);
    }
    setExpandedSections(newExpanded);
  };

  const formatNumber = (num: number) => {
    if (num >= 1e9) return `${(num / 1e9).toFixed(1)}B`;
    if (num >= 1e6) return `${(num / 1e6).toFixed(1)}M`;
    if (num >= 1e3) return `${(num / 1e3).toFixed(1)}K`;
    return num.toString();
  };

  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">
            <Brain className="w-6 h-6 text-purple-400" />
            Model Inspector
          </h2>
          <p className="text-gray-400 mt-1">
            Explore the complete architecture of the loaded model
          </p>
        </div>
        
        <div className={`flex items-center gap-2 px-3 py-1 rounded-full ${
          isConnected ? 'bg-green-900/30 text-green-400' : 'bg-red-900/30 text-red-400'
        }`}>
          <Activity className={`w-4 h-4 ${isConnected ? 'animate-pulse' : ''}`} />
          {isLoading ? 'Loading...' : (isConnected ? 'Model Connected' : 'Disconnected')}
        </div>
      </div>

      {/* Model Overview Section */}
      <div className="bg-gray-800 rounded-lg mb-4">
        <button
          onClick={() => toggleSection('overview')}
          className="w-full p-4 flex items-center justify-between hover:bg-gray-700/50 transition-colors rounded-lg"
        >
          <div className="flex items-center gap-3">
            {expandedSections.has('overview') ? 
              <ChevronDown className="w-5 h-5 text-gray-400" /> : 
              <ChevronRight className="w-5 h-5 text-gray-400" />
            }
            <Cpu className="w-5 h-5 text-blue-400" />
            <span className="font-semibold">Model Overview</span>
          </div>
          <span className="text-sm text-gray-400">{modelInfo.name}</span>
        </button>
        
        {expandedSections.has('overview') && (
          <div className="px-4 pb-4">
            <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">Total Parameters</div>
                <div className="text-2xl font-bold text-white">{formatNumber(modelInfo.totalParams)}</div>
                <div className="text-xs text-gray-500">
                  {modelInfo.totalParams > 1e9 ? `${(modelInfo.totalParams / 1e9).toFixed(1)} Billion` : 
                   modelInfo.totalParams > 1e6 ? `${(modelInfo.totalParams / 1e6).toFixed(1)} Million` : 
                   formatNumber(modelInfo.totalParams)}
                </div>
              </div>
              
              <div className="bg-gray-900 rounded-lg p-3">
                <div className="text-xs text-gray-400">Vocabulary Size</div>
                <div className="text-2xl font-bold text-purple-400">{formatNumber(modelInfo.vocabSize)}</div>
                <div className="text-xs text-gray-500">Unique tokens</div>
              </div>
              
              <div className="bg-gray-900 rounded-lg p-3">
                <div className="text-xs text-gray-400">Context Length</div>
                <div className="text-2xl font-bold text-green-400">{formatNumber(modelInfo.maxPositions)}</div>
                <div className="text-xs text-gray-500">Max tokens</div>
              </div>
              
              <div className="bg-gray-900 rounded-lg p-3">
                <div className="text-xs text-gray-400">Architecture</div>
                <div className="text-lg font-bold text-yellow-400">Transformer</div>
                <div className="text-xs text-gray-500">GPT-style</div>
              </div>
            </div>
            
            {/* Device Information */}
            {deviceInfo && (
              <div className="mt-4 p-3 bg-gray-900 rounded-lg">
                <div className="flex items-center gap-2 mb-2">
                  <Cpu className="w-4 h-4 text-blue-400" />
                  <span className="text-sm font-semibold">Device Information</span>
                </div>
                <div className="text-sm text-gray-400">
                  Running on: <span className="text-white font-mono">{deviceInfo}</span>
                </div>
              </div>
            )}
            
            {/* Model Configuration */}
            {modelConfig && (
              <div className="mt-4 p-3 bg-gray-900 rounded-lg">
                <div className="flex items-center gap-2 mb-2">
                  <Database className="w-4 h-4 text-purple-400" />
                  <span className="text-sm font-semibold">Configuration</span>
                </div>
                <div className="grid grid-cols-2 gap-2 text-xs">
                  <div className="flex justify-between">
                    <span className="text-gray-400">Activation:</span>
                    <span className="text-white font-mono">{String(modelConfig.activation_function)}</span>
                  </div>
                  <div className="flex justify-between">
                    <span className="text-gray-400">Cache:</span>
                    <span className="text-white font-mono">{modelConfig.use_cache ? 'Enabled' : 'Disabled'}</span>
                  </div>
                </div>
              </div>
            )}
          </div>
        )}
      </div>

      {/* Architecture Details Section */}
      <div className="bg-gray-800 rounded-lg mb-4">
        <button
          onClick={() => toggleSection('architecture')}
          className="w-full p-4 flex items-center justify-between hover:bg-gray-700/50 transition-colors rounded-lg"
        >
          <div className="flex items-center gap-3">
            {expandedSections.has('architecture') ? 
              <ChevronDown className="w-5 h-5 text-gray-400" /> : 
              <ChevronRight className="w-5 h-5 text-gray-400" />
            }
            <Layers className="w-5 h-5 text-green-400" />
            <span className="font-semibold">Architecture Details</span>
          </div>
          <span className="text-sm text-gray-400">{modelInfo.layers} layers × {modelInfo.heads} heads</span>
        </button>
        
        {expandedSections.has('architecture') && (
          <div className="px-4 pb-4">
            <div className="space-y-4">
              {/* Layer Structure */}
              <div className="bg-gray-900 rounded-lg p-4">
                <h4 className="text-sm font-semibold mb-3 flex items-center gap-2">
                  <GitBranch className="w-4 h-4 text-blue-400" />
                  Layer Structure (×{modelInfo.layers})
                </h4>
                <div className="space-y-2 text-sm">
                  <div className="flex items-center gap-2">
                    <div className="w-2 h-2 bg-green-400 rounded-full"></div>
                    <span className="text-gray-300">Multi-Head Attention</span>
                    <span className="text-gray-500">({modelInfo.heads} heads, {modelInfo.hiddenSize / modelInfo.heads} dims/head)</span>
                  </div>
                  <div className="flex items-center gap-2 ml-4">
                    <div className="w-1 h-1 bg-gray-400 rounded-full"></div>
                    <span className="text-gray-400">QKV Projection: {modelInfo.hiddenSize} → {modelInfo.hiddenSize * 3}</span>
                  </div>
                  <div className="flex items-center gap-2 ml-4">
                    <div className="w-1 h-1 bg-gray-400 rounded-full"></div>
                    <span className="text-gray-400">Output Projection: {modelInfo.hiddenSize} → {modelInfo.hiddenSize}</span>
                  </div>
                  
                  <div className="flex items-center gap-2 mt-2">
                    <div className="w-2 h-2 bg-purple-400 rounded-full"></div>
                    <span className="text-gray-300">Feed-Forward Network</span>
                    <span className="text-gray-500">(4× expansion)</span>
                  </div>
                  <div className="flex items-center gap-2 ml-4">
                    <div className="w-1 h-1 bg-gray-400 rounded-full"></div>
                    <span className="text-gray-400">FC1: {modelInfo.hiddenSize} → {modelInfo.hiddenSize * 4}</span>
                  </div>
                  <div className="flex items-center gap-2 ml-4">
                    <div className="w-1 h-1 bg-gray-400 rounded-full"></div>
                    <span className="text-gray-400">FC2: {modelInfo.hiddenSize * 4} → {modelInfo.hiddenSize}</span>
                  </div>
                  
                  <div className="flex items-center gap-2 mt-2">
                    <div className="w-2 h-2 bg-yellow-400 rounded-full"></div>
                    <span className="text-gray-300">Layer Normalization</span>
                  </div>
                  
                  <div className="flex items-center gap-2">
                    <div className="w-2 h-2 bg-red-400 rounded-full"></div>
                    <span className="text-gray-300">Residual Connections</span>
                  </div>
                </div>
              </div>

              {/* Data Flow */}
              <div className="bg-gray-900 rounded-lg p-4">
                <h4 className="text-sm font-semibold mb-3 flex items-center gap-2">
                  <Network className="w-4 h-4 text-purple-400" />
                  Data Flow Through Model
                </h4>
                <div className="font-mono text-xs text-gray-400 whitespace-pre">
{`Input Text
    ↓
[Token Embeddings] (${modelInfo.vocabSize.toLocaleString()} × ${modelInfo.hiddenSize.toLocaleString()})
    ↓
[+ Rotary Position Embeddings]
    ↓
╔═══════════════════════╗
║  Layer 0              ║
║  ├─ Attention (${modelInfo.heads}h)   ║
║  └─ FFN (${modelInfo.hiddenSize * 4}d)       ║
╚═══════════════════════╝
    ↓
    ... (${modelInfo.layers - 2} more layers)
    ↓
╔═══════════════════════╗
║  Layer ${modelInfo.layers - 1}             ║
║  ├─ Attention (${modelInfo.heads}h)   ║
║  └─ FFN (${modelInfo.hiddenSize * 4}d)       ║
╚═══════════════════════╝
    ↓
[Layer Norm]
    ↓
[Language Model Head]
    ↓
${modelInfo.vocabSize.toLocaleString()} Token Probabilities`}
                </div>
              </div>
            </div>
          </div>
        )}
      </div>

      {/* Accessible Components Section */}
      <div className="bg-gray-800 rounded-lg mb-4">
        <button
          onClick={() => toggleSection('accessible')}
          className="w-full p-4 flex items-center justify-between hover:bg-gray-700/50 transition-colors rounded-lg"
        >
          <div className="flex items-center gap-3">
            {expandedSections.has('accessible') ? 
              <ChevronDown className="w-5 h-5 text-gray-400" /> : 
              <ChevronRight className="w-5 h-5 text-gray-400" />
            }
            <Eye className="w-5 h-5 text-yellow-400" />
            <span className="font-semibold">What We Can Visualize</span>
          </div>
          <span className="text-sm text-gray-400">Everything is accessible</span>
        </button>
        
        {expandedSections.has('accessible') && (
          <div className="px-4 pb-4">
            <div className="grid grid-cols-1 md:grid-cols-2 gap-3">
              {modelInfo.accessible.map((item, idx) => (
                <div key={idx} className="flex items-start gap-2">
                  <Zap className="w-4 h-4 text-green-400 mt-0.5 flex-shrink-0" />
                  <span className="text-sm text-gray-300">{item}</span>
                </div>
              ))}
            </div>
            
            <div className="mt-4 p-4 bg-gray-900 rounded-lg border border-blue-900">
              <div className="flex items-start gap-2">
                <Info className="w-5 h-5 text-blue-400 flex-shrink-0 mt-0.5" />
                <div className="text-sm">
                  <p className="text-blue-400 font-semibold mb-1">Complete Transparency</p>
                  <p className="text-gray-400">
                    Every computation, weight, and decision in the model&apos;s {formatNumber(modelInfo.totalParams)} parameters 
                    is accessible. The &quot;black box&quot; becomes a &quot;glass box&quot; - we can visualize the 
                    entire thinking process as tokens flow through the network.
                  </p>
                </div>
              </div>
            </div>
          </div>
        )}
      </div>

      {/* Decision Path Visualization Section */}
      <div className="bg-gray-800 rounded-lg mb-4">
        <button
          onClick={() => toggleSection('decision-path')}
          className="w-full p-4 flex items-center justify-between hover:bg-gray-700/50 transition-colors rounded-lg"
        >
          <div className="flex items-center gap-3">
            {expandedSections.has('decision-path') ? 
              <ChevronDown className="w-5 h-5 text-gray-400" /> : 
              <ChevronRight className="w-5 h-5 text-gray-400" />
            }
            <GitBranch className="w-5 h-5 text-yellow-400" />
            <span className="font-semibold">Decision Path Analysis</span>
          </div>
          <span className="text-sm text-gray-400">Glass Box view</span>
        </button>
        
        {expandedSections.has('decision-path') && (
          <div className="px-4 pb-4">
            <Suspense fallback={
              <div className="h-[900px] bg-gray-900 rounded-lg flex items-center justify-center">
                <div className="text-gray-400">Loading decision path visualization...</div>
              </div>
            }>
              <DecisionPath3D />
            </Suspense>
          </div>
        )}
      </div>

      {/* 3D Visualization Section */}
      <div className="bg-gray-800 rounded-lg mb-4">
        <button
          onClick={() => toggleSection('3d')}
          className="w-full p-4 flex items-center justify-between hover:bg-gray-700/50 transition-colors rounded-lg"
        >
          <div className="flex items-center gap-3">
            {expandedSections.has('3d') ? 
              <ChevronDown className="w-5 h-5 text-gray-400" /> : 
              <ChevronRight className="w-5 h-5 text-gray-400" />
            }
            <Move3D className="w-5 h-5 text-blue-400" />
            <span className="font-semibold">3D Architecture Visualization</span>
          </div>
          <span className="text-sm text-gray-400">Interactive 3D model</span>
        </button>
        
        {expandedSections.has('3d') && (
          <div className="px-4 pb-4">
            <Suspense fallback={
              <div className="h-[800px] bg-gray-900 rounded-lg flex items-center justify-center">
                <div className="text-gray-400">Loading 3D visualization...</div>
              </div>
            }>
              <ModelArchitecture3D />
            </Suspense>
          </div>
        )}
      </div>

      {/* Computation Stats Section */}
      <div className="bg-gray-800 rounded-lg">
        <button
          onClick={() => toggleSection('computation')}
          className="w-full p-4 flex items-center justify-between hover:bg-gray-700/50 transition-colors rounded-lg"
        >
          <div className="flex items-center gap-3">
            {expandedSections.has('computation') ? 
              <ChevronDown className="w-5 h-5 text-gray-400" /> : 
              <ChevronRight className="w-5 h-5 text-gray-400" />
            }
            <Database className="w-5 h-5 text-red-400" />
            <span className="font-semibold">Computation Scale</span>
          </div>
          <span className="text-sm text-gray-400">Per token generation</span>
        </button>
        
        {expandedSections.has('computation') && (
          <div className="px-4 pb-4">
            <div className="space-y-3">
              <div className="flex items-center justify-between">
                <span className="text-sm text-gray-400">Operations per token</span>
                <span className="text-sm font-mono text-white">~{formatNumber(modelInfo.totalParams * 2)}</span>
              </div>
              <div className="flex items-center justify-between">
                <span className="text-sm text-gray-400">Attention computations</span>
                <span className="text-sm font-mono text-white">{modelInfo.layers * modelInfo.heads} heads</span>
              </div>
              <div className="flex items-center justify-between">
                <span className="text-sm text-gray-400">Probability calculations</span>
                <span className="text-sm font-mono text-white">{modelInfo.vocabSize.toLocaleString()} tokens</span>
              </div>
              <div className="flex items-center justify-between">
                <span className="text-sm text-gray-400">Memory footprint (FP32)</span>
                <span className="text-sm font-mono text-white">{((modelInfo.totalParams * 4) / (1024 * 1024 * 1024)).toFixed(2)} GB</span>
              </div>
              <div className="flex items-center justify-between">
                <span className="text-sm text-gray-400">Memory footprint (FP16)</span>
                <span className="text-sm font-mono text-white">{((modelInfo.totalParams * 2) / (1024 * 1024 * 1024)).toFixed(2)} GB</span>
              </div>
            </div>
            
            <div className="mt-4 p-3 bg-gray-900 rounded-lg">
              <p className="text-xs text-gray-500">
                Each token generation involves passing through all {modelInfo.layers} layers, 
                computing attention across {modelInfo.layers * modelInfo.heads} heads, and producing probabilities 
                for {modelInfo.vocabSize.toLocaleString()} possible next tokens.
              </p>
            </div>
          </div>
        )}
      </div>
    </div>
  );
}