frontend/src/pages/ClassificationPage.tsx

import { useState, useCallback, useEffect } from 'react';
import { Search, ChevronLeft, ChevronRight, FolderOpen, Upload, AlertTriangle } from 'lucide-react';
import { Panel } from '../components/Panel';
import { FileUpload } from '../components/FileUpload';
import { ResultsCard } from '../components/ResultsCard';
import { PreprocessingBadge } from '../components/PreprocessingBadge';
import { FeedbackSection } from '../components/FeedbackSection';
import { SessionHistory } from '../components/SessionHistory';
import { useImageContext } from '../lib/ImageContext';
import {
  classifyImage,
  getFilePreview,
  getFileType,
  isDicomFile,
  createSession,
  recordImageAnalyzed,
  type ClassificationResult,
  type PreprocessingInfo
} from '../lib/api';

interface ClassificationPageProps {
  onFeedbackUpdate?: () => void;
}

export function ClassificationPage({ onFeedbackUpdate }: ClassificationPageProps) {
  const imageContext = useImageContext();

  // Session state
  const [sessionId, setSessionId] = useState<string>('');
  const [feedbackRefresh, setFeedbackRefresh] = useState(0);

  // File state
  const [file, setFile] = useState<File | null>(null);
  const [preview, setPreview] = useState<string | null>(null);
  const [isLoadingPreview, setIsLoadingPreview] = useState(false);

  // Multiple files state
  const [files, setFiles] = useState<File[]>([]);
  const [currentIndex, setCurrentIndex] = useState(0);

  // Settings & results
  const [topK, setTopK] = useState(5);
  const [results, setResults] = useState<ClassificationResult[] | null>(null);
  const [preprocessingInfo, setPreprocessingInfo] = useState<PreprocessingInfo | null>(null);
  const [processedImage, setProcessedImage] = useState<string | null>(null);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);

  // Image view tab
  const [imageTab, setImageTab] = useState<'input' | 'processed'>('input');

  // Initialize session
  useEffect(() => {
    const initSession = async () => {
      try {
        const session = await createSession();
        setSessionId(session.session_id);
      } catch (err) {
        console.error('Failed to create session:', err);
      }
    };
    initSession();
  }, []);

  const loadPreview = useCallback(async (selectedFile: File) => {
    if (!isDicomFile(selectedFile.name)) {
      setPreview(URL.createObjectURL(selectedFile));
      return;
    }

    setIsLoadingPreview(true);
    try {
      const response = await getFilePreview(selectedFile);
      if (response.success) {
        setPreview(`data:image/png;base64,${response.preview}`);
      }
    } catch (err) {
      console.error('Failed to load DICOM preview:', err);
      setPreview(null);
    } finally {
      setIsLoadingPreview(false);
    }
  }, []);

  const handleSingleUpload = useCallback((uploadedFile: File) => {
    setFile(uploadedFile);
    setFiles([]);
    setCurrentIndex(0);
    setResults(null);
    setPreprocessingInfo(null);
    setProcessedImage(null);
    setError(null);
    setImageTab('input');
    loadPreview(uploadedFile);
  }, [loadPreview]);

  const handleFolderUpload = useCallback((e: React.ChangeEvent<HTMLInputElement>) => {
    const fileList = e.target.files;
    if (!fileList) return;

    const validFiles = Array.from(fileList).filter(f =>
      f.type.startsWith('image/') || isDicomFile(f.name)
    ).sort((a, b) => a.name.localeCompare(b.name));

    if (validFiles.length > 0) {
      setFiles(validFiles);
      setCurrentIndex(0);
      setFile(validFiles[0]);
      setResults(null);
      setPreprocessingInfo(null);
      setProcessedImage(null);
      setError(null);
      setImageTab('input');
      loadPreview(validFiles[0]);
    }
  }, [loadPreview]);

  const navigateImage = useCallback((direction: 'prev' | 'next') => {
    if (files.length === 0) return;

    let newIndex = currentIndex;
    if (direction === 'prev' && currentIndex > 0) {
      newIndex = currentIndex - 1;
    } else if (direction === 'next' && currentIndex < files.length - 1) {
      newIndex = currentIndex + 1;
    }

    if (newIndex !== currentIndex) {
      setCurrentIndex(newIndex);
      setFile(files[newIndex]);
      setResults(null);
      setPreprocessingInfo(null);
      setProcessedImage(null);
      setImageTab('input');
      loadPreview(files[newIndex]);
    }
  }, [files, currentIndex, loadPreview]);

  const handleClassify = async () => {
    if (!file) return;

    setIsLoading(true);
    setError(null);

    try {
      const response = await classifyImage(file, topK);
      setResults(response.predictions);
      setPreprocessingInfo(response.preprocessing);

      const processedImageData = response.preprocessing.processed_image_base64
        ? `data:image/png;base64,${response.preprocessing.processed_image_base64}`
        : null;

      if (processedImageData) {
        setProcessedImage(processedImageData);
      }
      setImageTab('processed');

      imageContext.setFile(file, preview);
      imageContext.setClassificationResults(response.predictions, processedImageData);

      if (sessionId) {
        await recordImageAnalyzed(sessionId);
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Classification failed');
      setResults(null);
      setPreprocessingInfo(null);
    } finally {
      setIsLoading(false);
    }
  };

  const handleFeedbackSubmitted = () => {
    setFeedbackRefresh(prev => prev + 1);
    onFeedbackUpdate?.();
  };

  const fileType = file ? getFileType(file.name) : null;
  const displayImage = imageTab === 'processed' && processedImage ? processedImage : preview;

  return (
    <div className="flex flex-1 min-h-0 overflow-hidden">
      {/* Left Panel - Image (60%) */}
      <div className="w-3/5 border-r border-dark-border bg-slate-900 flex flex-col min-h-0">
        {/* Compact Header */}
        <div className="flex-shrink-0 px-4 py-2 bg-slate-800 border-b border-slate-700 flex items-center justify-between">
          <div className="flex items-center gap-3">
            {/* Image Tab Toggle */}
            <div className="flex gap-1 bg-slate-700 p-0.5 rounded-lg">
              <button
                onClick={() => setImageTab('input')}
                className={`px-3 py-1 text-xs font-medium rounded-md transition-all ${imageTab === 'input'
                  ? 'bg-nvidia-green text-white'
                  : 'text-slate-400 hover:text-white'
                  }`}
              >
                Input
              </button>
              <button
                onClick={() => setImageTab('processed')}
                disabled={!processedImage}
                className={`px-3 py-1 text-xs font-medium rounded-md transition-all ${imageTab === 'processed'
                  ? 'bg-nvidia-green text-white'
                  : 'text-slate-400 hover:text-white disabled:opacity-40 disabled:cursor-not-allowed'
                  }`}
              >
                Processed
              </button>
            </div>

            {/* File info */}
            {file && (
              <span className="text-xs text-slate-400 truncate max-w-[150px]">
                {file.name}
              </span>
            )}

            {/* DICOM badge */}
            {fileType === 'dicom' && (
              <span className="px-2 py-0.5 bg-nvidia-green/20 text-nvidia-green text-xs rounded-full font-medium">
                DICOM
              </span>
            )}
          </div>

          {/* Folder navigation */}
          {files.length > 1 && (
            <div className="flex items-center gap-2">
              <button
                onClick={() => navigateImage('prev')}
                disabled={currentIndex === 0}
                className="p-1 text-slate-400 hover:text-white disabled:opacity-40"
              >
                <ChevronLeft className="w-4 h-4" />
              </button>
              <span className="text-xs text-slate-400">
                {currentIndex + 1}/{files.length}
              </span>
              <button
                onClick={() => navigateImage('next')}
                disabled={currentIndex === files.length - 1}
                className="p-1 text-slate-400 hover:text-white disabled:opacity-40"
              >
                <ChevronRight className="w-4 h-4" />
              </button>
            </div>
          )}
        </div>

        {/* Image Display - fills remaining space */}
        <div className="flex-1 min-h-0 p-4">
          {displayImage ? (
            <img
              src={displayImage}
              alt="Ultrasound"
              className="w-full h-full object-contain rounded-lg"
            />
          ) : (
            <FileUpload
              onUpload={handleSingleUpload}
              preview={null}
              currentFile={null}
              isLoading={isLoadingPreview}
            />
          )}
        </div>

        {/* Compact Control Bar */}
        <div className="flex-shrink-0 px-4 py-3 bg-slate-800 border-t border-slate-700">
          <div className="flex items-center gap-3">
            {/* Upload/Folder buttons */}
            <label className="flex items-center gap-1.5 px-3 py-1.5 bg-slate-700 hover:bg-slate-600 rounded-lg cursor-pointer transition-colors">
              <Upload className="w-3.5 h-3.5 text-slate-300" />
              <span className="text-xs text-slate-300 font-medium">Upload</span>
              <input
                type="file"
                accept="image/*,.dcm,.dicom"
                className="hidden"
                onChange={(e) => e.target.files?.[0] && handleSingleUpload(e.target.files[0])}
              />
            </label>

            <label className="flex items-center gap-1.5 px-3 py-1.5 bg-slate-700 hover:bg-slate-600 rounded-lg cursor-pointer transition-colors">
              <FolderOpen className="w-3.5 h-3.5 text-slate-300" />
              <span className="text-xs text-slate-300 font-medium">Folder</span>
              <input
                type="file"
                webkitdirectory=""
                directory=""
                multiple
                className="hidden"
                onChange={handleFolderUpload}
              />
            </label>

            <div className="w-px h-6 bg-slate-600" />

            {/* Top-K selector */}
            <div className="flex items-center gap-1.5">
              <span className="text-xs text-slate-400">Top</span>
              <select
                value={topK}
                onChange={(e) => setTopK(parseInt(e.target.value))}
                className="px-2 py-1 bg-slate-700 border border-slate-600 rounded text-xs text-white"
              >
                {[3, 5, 10, 13].map(k => (
                  <option key={k} value={k}>{k}</option>
                ))}
              </select>
            </div>

            <div className="flex-1" />

            {/* Classify Button */}
            <button
              onClick={handleClassify}
              disabled={!file || isLoading}
              className={`flex items-center gap-2 px-4 py-1.5 rounded-lg text-sm font-semibold transition-all ${!file
                ? 'bg-slate-600 text-slate-400 cursor-not-allowed'
                : 'bg-nvidia-green text-white hover:bg-nvidia-green-hover shadow-lg'
                }`}
            >
              {isLoading ? (
                <div className="w-4 h-4 border-2 border-white/30 border-t-white rounded-full animate-spin" />
              ) : (
                <Search className="w-4 h-4" />
              )}
              Classify
            </button>
          </div>

          {/* Error row */}
          {error && (
            <div className="mt-2 p-2 rounded-lg flex items-center gap-2 text-xs bg-red-500/10 text-red-400">
              <AlertTriangle className="w-3.5 h-3.5 flex-shrink-0" />
              {error}
            </div>
          )}
        </div>
      </div>

      {/* Right Panel - Results (40%) */}
      <div className="w-2/5 bg-white flex flex-col min-h-0">
        <Panel title="Results" className="flex-1 flex flex-col min-h-0">
          <div className="flex-1 overflow-y-auto space-y-4">
            {/* Preprocessing Badge */}
            {(fileType || preprocessingInfo) && (
              <PreprocessingBadge info={preprocessingInfo} fileType={fileType} />
            )}

            {/* Results Card */}
            <ResultsCard
              results={results}
              isLoading={isLoading}
            />

            {/* Feedback Section */}
            {results && results.length > 0 && file && (
              <FeedbackSection
                sessionId={sessionId}
                filename={file.name}
                fileType={fileType || 'image'}
                predictions={results}
                topPrediction={results[0]}
                preprocessedImageBase64={processedImage ? processedImage.split(',')[1] : undefined}
                onFeedbackSubmitted={handleFeedbackSubmitted}
                onViewCorrected={(correctedLabel) => imageContext.setCorrectedView(correctedLabel)}
              />
            )}

            {/* Session History */}
            {sessionId && (
              <SessionHistory
                sessionId={sessionId}
                refreshTrigger={feedbackRefresh}
              />
            )}
          </div>
        </Panel>
      </div>
    </div>
  );
}