Atomizer/atomizer-dashboard/frontend/src/pages/Home.tsx

import React, { useEffect, useState, useMemo } from 'react';
import { useNavigate } from 'react-router-dom';
import {
  Play,
  Pause,
  CheckCircle,
  Clock,
  RefreshCw,
  Zap,
  FileText,
  ChevronDown,
  ChevronUp,
  ChevronRight,
  Target,
  Activity,
  BarChart3,
  TrendingUp,
  ArrowRight,
  Folder,
  FolderOpen,
  Maximize2,
  X
} from 'lucide-react';
import { useStudy } from '../context/StudyContext';
import { Study } from '../types';
import { apiClient } from '../api/client';
import { MarkdownRenderer } from '../components/MarkdownRenderer';

const Home: React.FC = () => {
  const { studies, setSelectedStudy, refreshStudies, isLoading } = useStudy();
  const [selectedPreview, setSelectedPreview] = useState<Study | null>(null);
  const [readme, setReadme] = useState<string>('');
  const [readmeLoading, setReadmeLoading] = useState(false);
  const [sortField, setSortField] = useState<'name' | 'status' | 'trials' | 'bestValue'>('trials');
  const [sortDir, setSortDir] = useState<'asc' | 'desc'>('desc');
  const [expandedTopics, setExpandedTopics] = useState<Set<string>>(new Set());
  const [isFullscreen, setIsFullscreen] = useState(false);
  const navigate = useNavigate();

  // Group studies by topic, sorted by most recent first
  const studiesByTopic = useMemo(() => {
    const grouped: Record<string, Study[]> = {};
    studies.forEach(study => {
      const topic = study.topic || 'Other';
      if (!grouped[topic]) grouped[topic] = [];
      grouped[topic].push(study);
    });

    // Sort studies within each topic by last_modified (most recent first)
    Object.keys(grouped).forEach(topic => {
      grouped[topic].sort((a, b) => {
        const aTime = a.last_modified ? new Date(a.last_modified).getTime() : 0;
        const bTime = b.last_modified ? new Date(b.last_modified).getTime() : 0;
        return bTime - aTime; // Descending (most recent first)
      });
    });

    // Get most recent study time for each topic (for topic sorting)
    const topicMostRecent: Record<string, number> = {};
    Object.keys(grouped).forEach(topic => {
      const mostRecent = grouped[topic][0]?.last_modified;
      topicMostRecent[topic] = mostRecent ? new Date(mostRecent).getTime() : 0;
    });

    // Sort topics by most recent study (most recent first), 'Other' always last
    const sortedTopics = Object.keys(grouped).sort((a, b) => {
      if (a === 'Other') return 1;
      if (b === 'Other') return -1;
      return topicMostRecent[b] - topicMostRecent[a]; // Descending (most recent first)
    });

    const result: Record<string, Study[]> = {};
    sortedTopics.forEach(topic => {
      result[topic] = grouped[topic];
    });
    return result;
  }, [studies]);

  // Topics start collapsed by default - no initialization needed
  // Users can expand topics by clicking on them

  const toggleTopic = (topic: string) => {
    setExpandedTopics(prev => {
      const next = new Set(prev);
      if (next.has(topic)) {
        next.delete(topic);
      } else {
        next.add(topic);
      }
      return next;
    });
  };

  // Load README when a study is selected for preview
  useEffect(() => {
    if (selectedPreview) {
      loadReadme(selectedPreview.id);
    } else {
      setReadme('');
    }
  }, [selectedPreview]);

  const loadReadme = async (studyId: string) => {
    setReadmeLoading(true);
    try {
      const response = await apiClient.getStudyReadme(studyId);
      setReadme(response.content || 'No README found for this study.');
    } catch {
      setReadme('No README found for this study.');
    } finally {
      setReadmeLoading(false);
    }
  };

  const handleSelectStudy = (study: Study) => {
    setSelectedStudy(study);
    navigate('/dashboard');
  };

  const handleSort = (field: typeof sortField) => {
    if (sortField === field) {
      setSortDir(sortDir === 'asc' ? 'desc' : 'asc');
    } else {
      setSortField(field);
      setSortDir('desc');
    }
  };

  // Sort studies
  const sortedStudies = useMemo(() => {
    return [...studies].sort((a, b) => {
      let aVal: any, bVal: any;

      switch (sortField) {
        case 'name':
          aVal = (a.name || a.id).toLowerCase();
          bVal = (b.name || b.id).toLowerCase();
          break;
        case 'trials':
          aVal = a.progress.current;
          bVal = b.progress.current;
          break;
        case 'bestValue':
          aVal = a.best_value ?? Infinity;
          bVal = b.best_value ?? Infinity;
          break;
        case 'status':
          const statusOrder = { running: 0, paused: 1, completed: 2, not_started: 3 };
          aVal = statusOrder[a.status as keyof typeof statusOrder] ?? 4;
          bVal = statusOrder[b.status as keyof typeof statusOrder] ?? 4;
          break;
        default:
          return 0;
      }

      if (sortDir === 'asc') {
        return aVal < bVal ? -1 : aVal > bVal ? 1 : 0;
      } else {
        return aVal > bVal ? -1 : aVal < bVal ? 1 : 0;
      }
    });
  }, [studies, sortField, sortDir]);

  // Aggregate stats
  const aggregateStats = useMemo(() => {
    const totalStudies = studies.length;
    const runningStudies = studies.filter(s => s.status === 'running').length;
    const completedStudies = studies.filter(s => s.status === 'completed').length;
    const totalTrials = studies.reduce((sum, s) => sum + s.progress.current, 0);
    const studiesWithValues = studies.filter(s => s.best_value !== null);
    const bestOverall = studiesWithValues.length > 0
      ? studiesWithValues.reduce((best, curr) =>
          (curr.best_value! < best.best_value!) ? curr : best
        )
      : null;

    return { totalStudies, runningStudies, completedStudies, totalTrials, bestOverall };
  }, [studies]);

  const getStatusIcon = (status: string) => {
    switch (status) {
      case 'running':
        return <Play className="w-4 h-4 text-green-400" />;
      case 'paused':
        return <Pause className="w-4 h-4 text-orange-400" />;
      case 'completed':
        return <CheckCircle className="w-4 h-4 text-blue-400" />;
      default:
        return <Clock className="w-4 h-4 text-dark-400" />;
    }
  };

  const getStatusColor = (status: string) => {
    switch (status) {
      case 'running': return 'text-green-400 bg-green-500/10';
      case 'paused': return 'text-orange-400 bg-orange-500/10';
      case 'completed': return 'text-blue-400 bg-blue-500/10';
      default: return 'text-dark-400 bg-dark-600';
    }
  };

  return (
    <div className="min-h-screen bg-dark-900">
      {/* Header */}
      <header className="bg-dark-800/50 border-b border-dark-700 backdrop-blur-sm sticky top-0 z-10">
        <div className="max-w-[1920px] mx-auto px-6 py-4">
          <div className="flex items-center justify-between">
            <div className="flex items-center gap-4">
              <div className="w-11 h-11 bg-gradient-to-br from-primary-500 to-primary-700 rounded-xl flex items-center justify-center shadow-lg shadow-primary-500/20">
                <Zap className="w-6 h-6 text-white" />
              </div>
              <div>
                <h1 className="text-2xl font-bold text-white">Atomizer</h1>
                <p className="text-dark-400 text-sm">FEA Optimization Platform</p>
              </div>
            </div>
            <button
              onClick={() => refreshStudies()}
              disabled={isLoading}
              className="flex items-center gap-2 px-4 py-2 bg-dark-700 hover:bg-dark-600
                         text-white rounded-lg transition-all disabled:opacity-50 border border-dark-600"
            >
              <RefreshCw className={`w-4 h-4 ${isLoading ? 'animate-spin' : ''}`} />
              Refresh
            </button>
          </div>
        </div>
      </header>

      <main className="max-w-[1920px] mx-auto px-6 py-8">
        {/* Aggregate Stats Cards */}
        <div className="grid grid-cols-2 md:grid-cols-4 gap-4 mb-8">
          <div className="bg-dark-800 rounded-xl p-4 border border-dark-600">
            <div className="flex items-center gap-2 text-dark-400 text-sm mb-2">
              <BarChart3 className="w-4 h-4" />
              Total Studies
            </div>
            <div className="text-3xl font-bold text-white">{aggregateStats.totalStudies}</div>
          </div>

          <div className="bg-dark-800 rounded-xl p-4 border border-dark-600">
            <div className="flex items-center gap-2 text-green-400 text-sm mb-2">
              <Play className="w-4 h-4" />
              Running
            </div>
            <div className="text-3xl font-bold text-green-400">{aggregateStats.runningStudies}</div>
          </div>

          <div className="bg-dark-800 rounded-xl p-4 border border-dark-600">
            <div className="flex items-center gap-2 text-dark-400 text-sm mb-2">
              <Activity className="w-4 h-4" />
              Total Trials
            </div>
            <div className="text-3xl font-bold text-white">{aggregateStats.totalTrials.toLocaleString()}</div>
          </div>

          <div className="bg-dark-800 rounded-xl p-4 border border-dark-600">
            <div className="flex items-center gap-2 text-primary-400 text-sm mb-2">
              <Target className="w-4 h-4" />
              Best Overall
            </div>
            <div className="text-2xl font-bold text-primary-400">
              {aggregateStats.bestOverall?.best_value !== null && aggregateStats.bestOverall?.best_value !== undefined
                ? aggregateStats.bestOverall.best_value.toExponential(3)
                : 'N/A'}
            </div>
            {aggregateStats.bestOverall && (
              <div className="text-xs text-dark-400 mt-1 truncate">
                {aggregateStats.bestOverall.name || aggregateStats.bestOverall.id}
              </div>
            )}
          </div>
        </div>

        {/* Two-column layout: Table + Preview */}
        <div className="grid grid-cols-1 lg:grid-cols-2 gap-6">
          {/* Study Table */}
          <div className="bg-dark-800 rounded-xl border border-dark-600 overflow-hidden">
            <div className="px-6 py-4 border-b border-dark-600 flex items-center justify-between">
              <h2 className="text-lg font-semibold text-white flex items-center gap-2">
                <TrendingUp className="w-5 h-5 text-primary-400" />
                Studies
              </h2>
              <span className="text-sm text-dark-400">{studies.length} studies</span>
            </div>

            {isLoading ? (
              <div className="flex items-center justify-center py-16 text-dark-400">
                <RefreshCw className="w-6 h-6 animate-spin mr-3" />
                Loading studies...
              </div>
            ) : studies.length === 0 ? (
              <div className="text-center py-16 text-dark-400">
                <BarChart3 className="w-12 h-12 mx-auto mb-4 opacity-50" />
                <p>No studies found</p>
                <p className="text-sm mt-1 text-dark-500">Create a new study to get started</p>
              </div>
            ) : (
              <div className="max-h-[500px] overflow-y-auto">
                {Object.entries(studiesByTopic).map(([topic, topicStudies]) => {
                  const isExpanded = expandedTopics.has(topic);
                  const topicTrials = topicStudies.reduce((sum, s) => sum + s.progress.current, 0);
                  const runningCount = topicStudies.filter(s => s.status === 'running').length;

                  return (
                    <div key={topic} className="border-b border-dark-600 last:border-b-0">
                      {/* Topic Header */}
                      <button
                        onClick={() => toggleTopic(topic)}
                        className="w-full px-4 py-3 flex items-center justify-between hover:bg-dark-750 transition-colors"
                      >
                        <div className="flex items-center gap-3">
                          {isExpanded ? (
                            <FolderOpen className="w-5 h-5 text-primary-400" />
                          ) : (
                            <Folder className="w-5 h-5 text-dark-400" />
                          )}
                          <span className="text-white font-medium">{topic.replace(/_/g, ' ')}</span>
                          <span className="text-dark-500 text-sm">({topicStudies.length})</span>
                          {runningCount > 0 && (
                            <span className="flex items-center gap-1 text-xs text-green-400 bg-green-500/10 px-2 py-0.5 rounded-full">
                              <Play className="w-3 h-3" />
                              {runningCount} running
                            </span>
                          )}
                        </div>
                        <div className="flex items-center gap-4">
                          <span className="text-dark-400 text-sm">{topicTrials.toLocaleString()} trials</span>
                          {isExpanded ? (
                            <ChevronDown className="w-4 h-4 text-dark-400" />
                          ) : (
                            <ChevronRight className="w-4 h-4 text-dark-400" />
                          )}
                        </div>
                      </button>

                      {/* Topic Studies */}
                      {isExpanded && (
                        <div className="bg-dark-850">
                          {topicStudies.map((study) => {
                            const completionPercent = study.progress.total > 0
                              ? Math.round((study.progress.current / study.progress.total) * 100)
                              : 0;

                            return (
                              <div
                                key={study.id}
                                onClick={() => setSelectedPreview(study)}
                                className={`px-4 py-3 pl-12 flex items-center gap-4 border-t border-dark-700 hover:bg-dark-700 transition-colors cursor-pointer ${
                                  selectedPreview?.id === study.id ? 'bg-primary-900/20' : ''
                                }`}
                              >
                                {/* Study Name */}
                                <div className="flex-1 min-w-0">
                                  <span className="text-white font-medium truncate block">
                                    {study.name || study.id}
                                  </span>
                                  {study.name && (
                                    <span className="text-xs text-dark-500 truncate block">{study.id}</span>
                                  )}
                                </div>

                                {/* Status */}
                                <span className={`inline-flex items-center gap-1.5 px-2 py-1 rounded-full text-xs font-medium ${getStatusColor(study.status)}`}>
                                  {getStatusIcon(study.status)}
                                  {study.status}
                                </span>

                                {/* Progress */}
                                <div className="flex items-center gap-2 w-32">
                                  <div className="flex-1 h-2 bg-dark-600 rounded-full overflow-hidden">
                                    <div
                                      className={`h-full transition-all ${
                                        completionPercent >= 100 ? 'bg-green-500' :
                                        completionPercent >= 50 ? 'bg-primary-500' :
                                        'bg-yellow-500'
                                      }`}
                                      style={{ width: `${Math.min(completionPercent, 100)}%` }}
                                    />
                                  </div>
                                  <span className="text-dark-400 text-xs font-mono w-14 text-right">
                                    {study.progress.current}/{study.progress.total}
                                  </span>
                                </div>

                                {/* Best Value */}
                                <span className={`font-mono text-sm w-20 text-right ${study.best_value !== null ? 'text-primary-400' : 'text-dark-500'}`}>
                                  {study.best_value !== null ? study.best_value.toExponential(2) : 'N/A'}
                                </span>
                              </div>
                            );
                          })}
                        </div>
                      )}
                    </div>
                  );
                })}
              </div>
            )}
          </div>

          {/* Study Preview */}
          <div className="bg-dark-800 rounded-xl border border-dark-600 overflow-hidden flex flex-col">
            {selectedPreview ? (
              <>
                {/* Preview Header */}
                <div className="px-6 py-4 border-b border-dark-600 flex items-center justify-between">
                  <div className="flex items-center gap-3">
                    <div className="w-10 h-10 bg-dark-700 rounded-lg flex items-center justify-center">
                      <FileText className="w-5 h-5 text-primary-400" />
                    </div>
                    <div className="min-w-0">
                      <h2 className="text-lg font-semibold text-white truncate">
                        {selectedPreview.name || selectedPreview.id}
                      </h2>
                      <p className="text-dark-400 text-sm">Study Documentation</p>
                    </div>
                  </div>
                  <button
                    onClick={() => handleSelectStudy(selectedPreview)}
                    className="flex items-center gap-2 px-4 py-2 bg-primary-600 hover:bg-primary-500
                               text-white rounded-lg transition-all font-medium shadow-lg shadow-primary-500/20
                               hover:shadow-primary-500/30 whitespace-nowrap"
                  >
                    Open
                    <ArrowRight className="w-4 h-4" />
                  </button>
                </div>

                {/* Study Quick Stats */}
                <div className="px-6 py-3 border-b border-dark-600 flex items-center justify-between text-sm">
                  <div className="flex items-center gap-6">
                    <div className="flex items-center gap-2">
                      {getStatusIcon(selectedPreview.status)}
                      <span className="text-dark-300 capitalize">{selectedPreview.status}</span>
                    </div>
                    <div className="flex items-center gap-2 text-dark-400">
                      <Activity className="w-4 h-4" />
                      <span>{selectedPreview.progress.current} / {selectedPreview.progress.total} trials</span>
                    </div>
                    {selectedPreview.best_value !== null && (
                      <div className="flex items-center gap-2 text-primary-400">
                        <Target className="w-4 h-4" />
                        <span>Best: {selectedPreview.best_value.toExponential(4)}</span>
                      </div>
                    )}
                  </div>
                  <button
                    onClick={() => setIsFullscreen(true)}
                    className="p-2 hover:bg-dark-700 rounded-lg transition-colors text-dark-400 hover:text-white"
                    title="View fullscreen"
                  >
                    <Maximize2 className="w-4 h-4" />
                  </button>
                </div>

                {/* README Content */}
                <div className="flex-1 overflow-y-auto p-6">
                  {readmeLoading ? (
                    <div className="flex items-center justify-center py-16 text-dark-400">
                      <RefreshCw className="w-6 h-6 animate-spin mr-3" />
                      Loading documentation...
                    </div>
                  ) : (
                    <MarkdownRenderer content={readme} studyId={selectedPreview.id} />
                  )}
                </div>
              </>
            ) : (
              <div className="flex-1 flex items-center justify-center text-dark-400">
                <div className="text-center">
                  <FileText className="w-16 h-16 mx-auto mb-4 opacity-30" />
                  <p className="text-lg">Select a study to preview</p>
                  <p className="text-sm mt-1 text-dark-500">Click on any row in the table</p>
                </div>
              </div>
            )}
          </div>
        </div>
      </main>

      {/* Fullscreen README Modal */}
      {isFullscreen && selectedPreview && (
        <div className="fixed inset-0 z-50 bg-dark-900/95 backdrop-blur-sm overflow-hidden">
          <div className="h-full flex flex-col">
            {/* Modal Header */}
            <div className="flex-shrink-0 px-8 py-4 border-b border-dark-700 flex items-center justify-between bg-dark-800">
              <div className="flex items-center gap-4">
                <div className="w-10 h-10 bg-dark-700 rounded-lg flex items-center justify-center">
                  <FileText className="w-5 h-5 text-primary-400" />
                </div>
                <div>
                  <h2 className="text-xl font-semibold text-white">
                    {selectedPreview.name || selectedPreview.id}
                  </h2>
                  <p className="text-dark-400 text-sm">Study Documentation</p>
                </div>
              </div>
              <div className="flex items-center gap-3">
                <button
                  onClick={() => handleSelectStudy(selectedPreview)}
                  className="flex items-center gap-2 px-4 py-2 bg-primary-600 hover:bg-primary-500
                             text-white rounded-lg transition-all font-medium"
                >
                  Open Study
                  <ArrowRight className="w-4 h-4" />
                </button>
                <button
                  onClick={() => setIsFullscreen(false)}
                  className="p-2 hover:bg-dark-700 rounded-lg transition-colors text-dark-400 hover:text-white"
                  title="Close fullscreen"
                >
                  <X className="w-5 h-5" />
                </button>
              </div>
            </div>

            {/* Modal Content */}
            <div className="flex-1 overflow-y-auto p-8">
              <div className="max-w-4xl mx-auto">
                {readmeLoading ? (
                  <div className="flex items-center justify-center py-16 text-dark-400">
                    <RefreshCw className="w-6 h-6 animate-spin mr-3" />
                    Loading documentation...
                  </div>
                ) : (
                  <MarkdownRenderer content={readme} studyId={selectedPreview.id} />
                )}
              </div>
            </div>
          </div>
        </div>
      )}
    </div>
  );
};

export default Home;