app/page.tsx

'use client';

import { useState, useCallback, useMemo } from 'react';
import { Canvas } from '@react-three/fiber';
import { OrbitControls, Grid } from '@react-three/drei';
import { patterns, Pattern } from '@/lib/patterns';
import { OrigamiMesh } from '@/components/OrigamiMesh';
import { TrainingDashboard, TrainingEntry } from '@/components/TrainingDashboard';
import { parseFoldFile } from '@/lib/foldParser';

export default function Optigami() {
  const [selectedEntry, setSelectedEntry] = useState<TrainingEntry | null>(null);
  const [foldPercent, setFoldPercent] = useState(1);
  const [key, setKey] = useState(0);

  // Convert training entry's fold_data into a Pattern for the 3D viewer
  const activePattern: Pattern | null = useMemo(() => {
    if (!selectedEntry?.fold_data) return null;
    try {
      const parsed = parseFoldFile(JSON.stringify(selectedEntry.fold_data), selectedEntry.task_name);
      return parsed;
    } catch {
      return null;
    }
  }, [selectedEntry]);

  // Fallback to a default pattern when no training entry is selected
  const displayPattern = activePattern || patterns[0];

  const handleEntrySelect = useCallback((entry: TrainingEntry) => {
    setSelectedEntry(entry);
    setFoldPercent(1);
    setKey(k => k + 1);
  }, []);

  return (
    <div className="flex h-screen w-full bg-zinc-950 text-zinc-100 font-sans overflow-hidden">
      {/* Left panel — Training Dashboard */}
      <div className="w-96 flex-shrink-0 border-r border-zinc-800 bg-zinc-900 flex flex-col">
        <div className="px-5 py-4 border-b border-zinc-800 flex items-center justify-between">
          <div>
            <h1 className="text-lg font-semibold tracking-tight">Optigami</h1>
            <p className="text-[11px] text-zinc-500">RL Training Environment</p>
          </div>
          <div className="flex items-center gap-2">
            <a
              href="https://huggingface.co/spaces/openenv-community/optigami_"
              target="_blank"
              rel="noopener noreferrer"
              className="text-[10px] text-zinc-500 hover:text-zinc-300 bg-zinc-800 px-2 py-1 rounded border border-zinc-700"
            >
              OpenEnv 0.2.1
            </a>
          </div>
        </div>

        <div className="flex-1 overflow-y-auto p-4">
          <TrainingDashboard onEntrySelect={handleEntrySelect} />
        </div>
      </div>

      {/* Right side — 3D viewer + detail */}
      <div className="flex-1 flex flex-col">
        {/* Top bar with context about selected entry */}
        {selectedEntry && (
          <div className="flex-shrink-0 border-b border-zinc-800 bg-zinc-900/50 px-6 py-3 flex items-center gap-6 text-xs">
            <div>
              <span className="text-zinc-500">Step</span>{' '}
              <span className="font-mono text-zinc-200">#{selectedEntry.step}</span>
            </div>
            <div>
              <span className="text-zinc-500">Task</span>{' '}
              <span className="text-zinc-200">{selectedEntry.task_name}</span>
            </div>
            <div>
              <span className="text-zinc-500">Reward</span>{' '}
              <span className={`font-mono font-semibold ${
                selectedEntry.reward >= 15 ? 'text-green-400' :
                selectedEntry.reward >= 5 ? 'text-yellow-400' :
                selectedEntry.reward >= 0 ? 'text-orange-400' : 'text-red-400'
              }`}>{selectedEntry.reward.toFixed(2)}</span>
            </div>
            <div>
              <span className="text-zinc-500">Similarity</span>{' '}
              <span className="font-mono text-indigo-400">
                {(selectedEntry.shape_similarity * 100).toFixed(1)}%
              </span>
            </div>
            {selectedEntry.error && (
              <div className="text-red-400/80 truncate flex-1">{selectedEntry.error}</div>
            )}
            <div className="ml-auto flex items-center gap-2">
              <span className="text-zinc-500">Fold</span>
              <input
                type="range"
                min="0"
                max="1"
                step="0.01"
                value={foldPercent}
                onChange={(e) => setFoldPercent(parseFloat(e.target.value))}
                className="w-24 accent-indigo-500"
              />
              <span className="font-mono text-zinc-400 w-8">{Math.round(foldPercent * 100)}%</span>
            </div>
          </div>
        )}

        {/* 3D Canvas */}
        <div className="flex-1 relative">
          {!selectedEntry && (
            <div className="absolute inset-0 flex items-center justify-center z-10 pointer-events-none">
              <div className="text-center">
                <div className="text-zinc-500 text-sm mb-1">Waiting for training data</div>
                <div className="text-zinc-600 text-xs">
                  Start a GRPO training run in the Colab notebook
                </div>
              </div>
            </div>
          )}
          <Canvas camera={{ position: [0, 0, 3], fov: 45 }}>
            <ambientLight intensity={0.5} />
            <directionalLight position={[5, 5, 5]} intensity={1} castShadow />
            <directionalLight position={[-5, -5, -5]} intensity={0.2} />

            <group key={key}>
              <OrigamiMesh pattern={displayPattern} foldPercent={foldPercent} />
            </group>

            <OrbitControls makeDefault />
            <Grid
              infiniteGrid
              fadeDistance={10}
              sectionColor="#333"
              cellColor="#222"
              position={[0, 0, -0.01]}
            />
          </Canvas>

          {/* 2D crease pattern overlay */}
          {activePattern && (
            <div className="absolute bottom-4 left-4 w-40 h-40 bg-zinc-900/90 rounded-lg border border-zinc-700/50 p-2 backdrop-blur-sm">
              <div className="text-[9px] uppercase text-zinc-500 tracking-wider mb-1">Crease Pattern</div>
              <svg viewBox="-1.2 -1.2 2.4 2.4" className="w-full h-full">
                <g transform="scale(1, -1)">
                  {activePattern.faces.map((face, i) => {
                    const v1 = activePattern.vertices[face[0]];
                    const v2 = activePattern.vertices[face[1]];
                    const v3 = activePattern.vertices[face[2]];
                    return (
                      <polygon
                        key={`f-${i}`}
                        points={`${v1[0]},${v1[1]} ${v2[0]},${v2[1]} ${v3[0]},${v3[1]}`}
                        fill="#3f3f46"
                        stroke="#52525b"
                        strokeWidth="0.01"
                      />
                    );
                  })}
                  {activePattern.creases.map((crease, i) => {
                    const v1 = activePattern.vertices[crease.edge[0]];
                    const v2 = activePattern.vertices[crease.edge[1]];
                    const color = crease.type === 'mountain' ? '#ef4444' : '#3b82f6';
                    return (
                      <line
                        key={`c-${i}`}
                        x1={v1[0]} y1={v1[1]}
                        x2={v2[0]} y2={v2[1]}
                        stroke={color}
                        strokeWidth="0.03"
                        strokeLinecap="round"
                      />
                    );
                  })}
                </g>
              </svg>
            </div>
          )}
        </div>
      </div>
    </div>
  );
}