frontend/src/components/viewer/PathoraViewer.tsx

import { useEffect, useRef, useState } from "react";
import OpenSeadragon from "openseadragon";
import { ToolsSidebar } from "./ToolsSidebar";
import { TopToolbar } from "./TopToolbar";
import { AnnotationCanvas, type Annotation } from "./AnnotationCanvas";
import { ArrowLeft } from "lucide-react";
import { useNavigate } from "react-router-dom";
import "./viewer.css";

interface PathoraViewerProps {
  imageUrl?: string;
  slideName?: string;
}

export type Tool = "none" | "select" | "rectangle" | "polygon" | "ellipse" | "brush";
type UploadedSlide = {
  id: string;
  name: string;
  uploadedAt: string;
  levelCount: number;
  levelDimensions: number[][];
};

export function PathoraViewer({ 

  imageUrl = "", 

  slideName = "Pathora Viewer" 

}: PathoraViewerProps) {
  console.log("PathoraViewer component rendering with:", { imageUrl, slideName });
  
  const viewerRef = useRef<HTMLDivElement>(null);
  const osdViewerRef = useRef<OpenSeadragon.Viewer | null>(null);
  const navigate = useNavigate();
  const [selectedTool, setSelectedTool] = useState<Tool>("none");
  const [zoomLevel, setZoomLevel] = useState<number>(1);
  const [showAnnotations, setShowAnnotations] = useState(true);
  const [showHeatmap, setShowHeatmap] = useState(false);
  const [annotations, setAnnotations] = useState<Annotation[]>([]);
  const [selectedAnnotationId, setSelectedAnnotationId] = useState<string | null>(null);
  const [activeLabel, setActiveLabel] = useState("Tumor");
  const [isSidebarCollapsed, setIsSidebarCollapsed] = useState(false);
  const [uploadedSlides, setUploadedSlides] = useState<UploadedSlide[]>([]);
  const [showOriginal, setShowOriginal] = useState(true);
  const [showHematoxylin, setShowHematoxylin] = useState(false);
  const [showEosin, setShowEosin] = useState(false);
  const [isLoading, setIsLoading] = useState(true);
  const [loadError, setLoadError] = useState<string | null>(null);
  const [tileServerUrl, setTileServerUrl] = useState("http://localhost:8001");
  const [slideId, setSlideId] = useState("slide-1");
  const [tileSize, setTileSize] = useState(256);
  const [tileMode, setTileMode] = useState<"none" | "image" | "tiles">("none");
  const [slideFile, setSlideFile] = useState<File | null>(null);
  const [autoSlideId, setAutoSlideId] = useState(true);
  const [tileMeta, setTileMeta] = useState<{
    width: number;
    height: number;
    level_count: number;
    level_dimensions: number[][];
    level_downsamples: number[];
    mpp_x?: number | null;
    mpp_y?: number | null;
  } | null>(null);
  const [tileLoadError, setTileLoadError] = useState<string | null>(null);
  const [isTileLoading, setIsTileLoading] = useState(false);
  const channelItemsRef = useRef<{
    original?: OpenSeadragon.TiledImage;
    hematoxylin?: OpenSeadragon.TiledImage;
    eosin?: OpenSeadragon.TiledImage;
  }>({});

  const normalizeBaseUrl = (value: string) => {
    return value.replace(/\/$/, "");
  };

  const autoSlideIdLabel = autoSlideId ? "Auto ID enabled" : "Manual ID";

  const generateSlideId = () => {
    const now = new Date();
    const stamp = now
      .toISOString()
      .replace(/[-:]/g, "")
      .replace("T", "-")
      .slice(0, 15);
    const rand = Math.random().toString(36).slice(2, 6);
    return `slide-${stamp}-${rand}`;
  };

  const buildTileSource = (meta: {

    width: number;

    height: number;

    level_count: number;

    level_downsamples: number[];

  }, channel: "original" | "hematoxylin" | "eosin" = "original") => {
    const baseUrl = normalizeBaseUrl(tileServerUrl);
    const maxLevel = Math.max(0, meta.level_count - 1);

    const tileSource = new OpenSeadragon.TileSource({
      width: meta.width,
      height: meta.height,
      tileSize,
      minLevel: 0,
      maxLevel,
    });

    tileSource.getLevelScale = (level: number) => {
      const slideLevel = Math.max(0, meta.level_count - 1 - level);
      const downsample = meta.level_downsamples[slideLevel] || 1;
      return 1 / downsample;
    };

    tileSource.getTileUrl = (level: number, x: number, y: number) => {
      const slideLevel = Math.max(0, meta.level_count - 1 - level);
      return `${baseUrl}/tiles/${slideId}/${slideLevel}/${x}/${y}.jpg?tile_size=${tileSize}&channel=${channel}`;
    };

    return tileSource;
  };

  const addUploadedSlide = (

    id: string,

    name: string,

    levelCount: number,

    levelDimensions: number[][]

  ) => {
    const uploadedAt = new Date().toLocaleString();
    setUploadedSlides((prev) => {
      const filtered = prev.filter((slide) => slide.id !== id);
      return [{ id, name, uploadedAt, levelCount, levelDimensions }, ...filtered].slice(0, 20);
    });
  };

  const handleLoadSlide = async () => {
    if (!slideId.trim()) {
      setTileLoadError("Slide id is required.");
      return;
    }

    if (!slideFile) {
      setTileLoadError("Please choose a WSI file to upload.");
      return;
    }

    const baseUrl = normalizeBaseUrl(tileServerUrl);
    setIsTileLoading(true);
    setTileLoadError(null);

    try {
      const formData = new FormData();
      formData.append("file", slideFile);

      const uploadRes = await fetch(`${baseUrl}/slides/${slideId}/upload`, {
        method: "POST",
        body: formData,
      });

      if (!uploadRes.ok) {
        const text = await uploadRes.text();
        throw new Error(text || "Failed to upload slide");
      }

      const metaRes = await fetch(`${baseUrl}/slides/${slideId}/metadata`);
      if (!metaRes.ok) {
        const text = await metaRes.text();
        throw new Error(text || "Failed to read metadata");
      }

      const meta = await metaRes.json();
      setTileMeta(meta);
      setTileMode("tiles");
      addUploadedSlide(
        slideId,
        slideFile?.name || "Untitled slide",
        meta.level_count || 1,
        meta.level_dimensions || []
      );
    } catch (error: any) {
      setTileLoadError(error?.message || "Failed to upload slide");
    } finally {
      setIsTileLoading(false);
    }
  };

  const handleSelectUploadedSlide = async (id: string) => {
    const baseUrl = normalizeBaseUrl(tileServerUrl);
    setSlideId(id);
    setTileLoadError(null);
    setIsTileLoading(true);

    const fetchMeta = async () => {
      const metaRes = await fetch(`${baseUrl}/slides/${id}/metadata`);
      if (!metaRes.ok) {
        const text = await metaRes.text();
        const error = new Error(text || "Failed to read metadata") as Error & {
          status?: number;
        };
        error.status = metaRes.status;
        throw error;
      }
      return metaRes.json();
    };

    try {
      let meta: any;
      try {
        meta = await fetchMeta();
      } catch (error: any) {
        if (error?.status === 404) {
          const reloadRes = await fetch(`${baseUrl}/slides/${id}/reload`, { method: "POST" });
          if (!reloadRes.ok) {
            const text = await reloadRes.text();
            throw new Error(text || "Failed to reload slide");
          }
          meta = await fetchMeta();
        } else {
          throw error;
        }
      }

      setTileMeta(meta);
      setTileMode("tiles");
    } catch (error: any) {
      setTileLoadError(error?.message || "Failed to load slide");
    } finally {
      setIsTileLoading(false);
    }
  };

  // Initialize OpenSeadragon viewer
  useEffect(() => {
    if (!viewerRef.current || osdViewerRef.current) return;

    console.log("Initializing OpenSeadragon with image:", imageUrl);
    console.log("Viewer ref:", viewerRef.current);

    try {
      const viewer = OpenSeadragon({
        element: viewerRef.current,
        prefixUrl: "https://cdnjs.cloudflare.com/ajax/libs/openseadragon/4.1.0/images/",
        crossOriginPolicy: "Anonymous",
        showNavigator: true,
        navigatorPosition: "BOTTOM_RIGHT",
        navigatorSizeRatio: 0.15,
        showNavigationControl: false,
        minZoomImageRatio: 0.5,
        maxZoomPixelRatio: 3,
        visibilityRatio: 0.5,
        constrainDuringPan: true,
        animationTime: 0.5,
        gestureSettingsMouse: {
          clickToZoom: false,
          dblClickToZoom: true,
        },
      });

      console.log("OpenSeadragon viewer created:", viewer);

      // Add error handler
      viewer.addHandler("open-failed", (event: any) => {
        console.error("Failed to open image:", event);
        setLoadError("Failed to load image. Please check the image path.");
        setIsLoading(false);
      });

      // Add success handler
      viewer.addHandler("open", () => {
        console.log("Image loaded successfully");
        setIsLoading(false);
        setLoadError(null);
      });

      // Update zoom level on zoom
      viewer.addHandler("zoom", () => {
        const zoom = viewer.viewport.getZoom();
        setZoomLevel(zoom);
      });

      osdViewerRef.current = viewer;

      return () => {
        if (osdViewerRef.current) {
          osdViewerRef.current.destroy();
          osdViewerRef.current = null;
        }
      };
    } catch (error) {
      console.error("Error initializing OpenSeadragon:", error);
    }
  }, [imageUrl]);

  useEffect(() => {
    if (!osdViewerRef.current) return;

    const viewer = osdViewerRef.current;

    if (tileMode === "tiles" && tileMeta) {
      setIsLoading(true);
      setLoadError(null);
      viewer.open(buildTileSource(tileMeta, "original"));
      return;
    }

    if (tileMode === "image" && imageUrl) {
      setIsLoading(true);
      setLoadError(null);
      viewer.open({
        type: "image",
        url: imageUrl,
      });
      return;
    }

    if (tileMode === "none") {
      setIsLoading(false);
      setLoadError(null);
      viewer.close();
    }
  }, [imageUrl, tileMeta, tileMode, tileServerUrl, tileSize, slideId]);

  useEffect(() => {
    if (!osdViewerRef.current || !tileMeta || tileMode !== "tiles") return;
    const viewer = osdViewerRef.current;

    const originalItem = viewer.world.getItemAt(0);
    if (originalItem) {
      channelItemsRef.current.original = originalItem;
      originalItem.setOpacity(showOriginal ? 1 : 0);
    }

    const ensureChannel = (

      key: "hematoxylin" | "eosin",

      enabled: boolean,

      channel: "hematoxylin" | "eosin"

    ) => {
      const existing = channelItemsRef.current[key];
      if (existing) {
        existing.setOpacity(enabled ? 1 : 0);
        return;
      }
      if (!enabled) return;

      viewer.addTiledImage({
        tileSource: buildTileSource(tileMeta, channel),
        opacity: 1,
        success: (event: any) => {
          channelItemsRef.current[key] = event.item as OpenSeadragon.TiledImage;
        },
      });
    };

    ensureChannel("hematoxylin", showHematoxylin, "hematoxylin");
    ensureChannel("eosin", showEosin, "eosin");
  }, [tileMeta, tileMode, showOriginal, showHematoxylin, showEosin, tileServerUrl, tileSize, slideId]);

  const handleAnnotationComplete = (annotation: Annotation) => {
    console.log("Annotation completed:", annotation);
    setAnnotations((prev) => [...prev, annotation]);
    // Auto-select the newly created annotation
    setSelectedAnnotationId(annotation.id);
  };

  const handleAnnotationSelected = (annotationId: string | null) => {
    setSelectedAnnotationId(annotationId);
  };

  const handleUndo = () => {
    if (annotations.length > 0) {
      const newAnnotations = annotations.slice(0, -1);
      setAnnotations(newAnnotations);
      setSelectedAnnotationId(null);
    }
  };

  const handleDelete = () => {
    if (selectedAnnotationId) {
      const newAnnotations = annotations.filter(
        (ann) => ann.id !== selectedAnnotationId
      );
      setAnnotations(newAnnotations);
      setSelectedAnnotationId(null);
    }
  };

  const handleZoomIn = () => {
    if (osdViewerRef.current) {
      const currentZoom = osdViewerRef.current.viewport.getZoom();
      osdViewerRef.current.viewport.zoomTo(currentZoom * 1.2);
    }
  };

  const handleZoomOut = () => {
    if (osdViewerRef.current) {
      const currentZoom = osdViewerRef.current.viewport.getZoom();
      osdViewerRef.current.viewport.zoomTo(currentZoom / 1.2);
    }
  };

  const zoomPresets = [1, 5, 10, 20, 40];
  const micronsPerPixel = tileMeta?.mpp_x ?? tileMeta?.mpp_y ?? null;
  const imageMeta = {
    stain: "H&E",
    width: tileMeta?.width ?? null,
    height: tileMeta?.height ?? null,
    levelCount: tileMeta?.level_count ?? null,
    mpp: micronsPerPixel,
    slideId,
  };
  const isTileLoaded = tileMode === "tiles" && !!tileMeta;

  const handleZoomPreset = (level: number) => {
    if (osdViewerRef.current) {
      osdViewerRef.current.viewport.zoomTo(level);
    }
  };

  return (
    <div className="flex flex-col h-screen bg-gray-100">

      {/* Header with back button */}

      <header className="h-16 bg-gradient-to-r from-teal-700 to-teal-600 text-white flex items-center px-6 shadow-md">

        <button

          onClick={() => navigate("/")}

          className="flex items-center space-x-2 hover:bg-teal-800/50 px-3 py-2 rounded-lg transition-colors"

        >

          <ArrowLeft className="w-5 h-5" />

          <span className="font-medium">Back to Analysis</span>

        </button>

        <div className="flex-1 text-center">

          <h1 className="text-2xl font-bold">Pathora Viewer</h1>

          <p className="text-xs text-teal-100">Advanced Whole Slide Imaging Platform</p>

        </div>

        <div className="w-40"></div> {/* Spacer for centering */}

      </header>



      <div className="flex flex-1 overflow-hidden">

        {/* Tools Sidebar */}

        <ToolsSidebar

          selectedTool={selectedTool}

          onToolChange={setSelectedTool}

          annotations={annotations}

          selectedAnnotationId={selectedAnnotationId}

          onSelectAnnotation={handleAnnotationSelected}

          uploadedSlides={uploadedSlides}

          tileServerUrl={tileServerUrl}

          onTileServerUrlChange={setTileServerUrl}

          onSlideFileChange={(file) => {

            setSlideFile(file);

            if (file && autoSlideId) {

              setSlideId(generateSlideId());

            }

          }}

          slideFileName={slideFile?.name ?? null}

          onUploadSlide={handleLoadSlide}

          isTileLoading={isTileLoading}

          tileLoadError={tileLoadError}

          onSelectUploadedSlide={handleSelectUploadedSlide}

          activeLabel={activeLabel}

          onLabelChange={setActiveLabel}

          imageMeta={imageMeta}

          channelVisibility={{

            original: showOriginal,

            hematoxylin: showHematoxylin,

            eosin: showEosin,

          }}

          onChannelToggle={(channel, value) => {

            if (channel === "original") setShowOriginal(value);

            if (channel === "hematoxylin") setShowHematoxylin(value);

            if (channel === "eosin") setShowEosin(value);

          }}

          isTileLoaded={isTileLoaded}

          isCollapsed={isSidebarCollapsed}

          onToggleCollapsed={() => setIsSidebarCollapsed((prev) => !prev)}

        />



        {/* Main Viewer Area */}

        <div className="flex-1 flex flex-col">{/* Top Toolbar */}

          <TopToolbar

            slideName={slideName}

            zoomLevel={zoomLevel}

            zoomPresets={zoomPresets}

            onZoomPreset={handleZoomPreset}

            micronsPerPixel={micronsPerPixel}

            showAnnotations={showAnnotations}

            showHeatmap={showHeatmap}

            onUndo={handleUndo}

            onDelete={handleDelete}

            onZoomIn={handleZoomIn}

            onZoomOut={handleZoomOut}

            onToggleAnnotations={() => setShowAnnotations(!showAnnotations)}

            onToggleHeatmap={() => setShowHeatmap(!showHeatmap)}

            canUndo={annotations.length > 0}

            canDelete={selectedAnnotationId !== null}

          />



          <div className="flex-1 flex overflow-hidden">

            {/* Viewer Container */}

            <div className="flex-1 relative">

              <div

                ref={viewerRef}

                className="absolute inset-0 bg-black"

                style={{ width: "100%", height: "100%" }}

              />



              {/* Loading State */}

              {isLoading && (

                <div className="absolute inset-0 flex items-center justify-center bg-gray-900 bg-opacity-75 z-50">

                  <div className="text-center">

                    <div className="inline-block animate-spin rounded-full h-16 w-16 border-t-4 border-b-4 border-teal-500 mb-4"></div>

                    <p className="text-white text-lg font-semibold">Loading slide viewer...</p>

                    <p className="text-gray-300 text-sm mt-2">Initializing OpenSeadragon</p>

                  </div>

                </div>

              )}



              {/* Error State */}

              {loadError && (

                <div className="absolute inset-0 flex items-center justify-center bg-gray-900 bg-opacity-75 z-50">

                  <div className="bg-white rounded-lg p-6 max-w-md text-center">

                    <div className="text-red-500 text-5xl mb-4">⚠️</div>

                    <h3 className="text-xl font-bold text-gray-800 mb-2">Failed to Load Image</h3>

                    <p className="text-gray-600 mb-4">{loadError}</p>

                    <button

                      onClick={() => window.location.reload()}

                      className="px-4 py-2 bg-teal-600 text-white rounded-lg hover:bg-teal-700 transition-colors"

                    >

                      Retry

                    </button>

                  </div>

                </div>

              )}



              {/* Annotation Canvas */}

              <AnnotationCanvas

                viewer={osdViewerRef.current}

                tool={selectedTool}

                onAnnotationComplete={handleAnnotationComplete}

                activeLabel={activeLabel}

                onAnnotationSelected={handleAnnotationSelected}

                annotations={annotations}

                selectedAnnotationId={selectedAnnotationId}

                showAnnotations={showAnnotations}

              />

            



            {/* Annotations count */}

            {showAnnotations && annotations.length > 0 && (

              <div className="absolute bottom-4 left-4 bg-white px-3 py-2 rounded shadow-md text-sm z-40">

                <span className="font-semibold">Annotations:</span> {annotations.length}

              </div>

            )}



            {/* Heatmap overlay placeholder */}

            {showHeatmap && (

              <div className="absolute inset-0 pointer-events-none bg-gradient-to-br from-red-500/20 via-yellow-500/20 to-green-500/20" />

            )}



              {tileMode === "none" && !isLoading && !loadError && (

                <div className="absolute inset-0 flex items-center justify-center bg-gray-900/60 z-40">

                  <div className="rounded-lg px-5 py-4 text-center">

                    <div className="text-sm font-semibold text-white">Upload a WSI to start</div>

                    <div className="text-xs text-white/80 mt-1">

                      Use the uploader in the left Uploads tab to load a slide.

                    </div>

                  </div>

                </div>

              )}

            </div>

          </div>

        </div>

      </div>

    </div>
  );
}