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>
  );
}