'use client';

import { StructRowProxy } from 'apache-arrow';
import clsx from 'clsx';
import { useCallback, useEffect, useRef, useState } from 'react';

import { useStores } from '@/app/(root)/AppProviders';
import Button, { ButtonVariant } from '@/components/button/Button';
import SpinnerSVG from '@/components/svg/SpinnerSVG';
import { TextInput } from '@/components/textarea/TextInput';
import { SearchIcon } from '@/icons/generated';
import { useApiClient } from '@/lib/apiClient';
import { staticAssetUrl } from '@/utils/staticAssetUrl';

import { Highlights } from './Highlights';
import SongTooltip from './SongTooltip';
import { Scatterplot } from './deepscatter/src/scatterplot';
import { Qid } from './deepscatter/src/tixrixqid';
import { uint8ArrayToUuid } from './utils';

// Global singleton with proper cancellation to prevent multiple instances
let globalScatterplot: Scatterplot | null = null;
let globalAbortController: AbortController | null = null;

const isZoomReady = () => {
  try {
    return !!globalScatterplot?.zoom;
  } catch {
    return false;
  }
};

const LABEL_OPTIONS = {
  font: 'Neue Montreal',
  useColorScale: true,
};

export default function SunoverseClientV4() {
  const containerRef = useRef<HTMLDivElement>(null);
  const scatterplotRef = useRef<Scatterplot | null>(null);
  const hoveredPointQidRef = useRef<Qid | null>(null);
  const [genreSearchText, setGenreSearchText] = useState('');
  const [genreLabels, setGenreLabels] = useState<any[]>([]);

  // Loading states
  const [scatterplotReady, setScatterplotReady] = useState(false);
  const [isLocatingGenre, setIsLocatingGenre] = useState(false);
  const [isLocatingMyself, setIsLocatingMyself] = useState(false);

  const { clips, playbar, session } = useStores();
  const clipPromises = useRef<Map<string, Promise<any>>>(new Map());

  const [tooltipData, setTooltipData] = useState<any>(null);
  const [tooltipPosition, setTooltipPosition] = useState<[number, number]>([
    0, 0,
  ]);

  // Use refs to avoid stale closures without causing re-runs
  const clipsRef = useRef(clips);
  const playbarRef = useRef(playbar);
  const setTooltipDataRef = useRef(setTooltipData);
  const setTooltipPositionRef = useRef(setTooltipPosition);

  const highlightsRef = useRef<Highlights | null>(null);

  // Update refs on every render
  clipsRef.current = clips;
  playbarRef.current = playbar;
  setTooltipDataRef.current = setTooltipData;
  setTooltipPositionRef.current = setTooltipPosition;

  // API things
  const apiClient = useApiClient();

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

    const initializeScatterplot = async () => {
      globalAbortController?.abort();
      globalAbortController = new AbortController();
      const controller = globalAbortController;

      if (globalScatterplot) {
        globalScatterplot.destroy();
        globalScatterplot = null;
      }
      setScatterplotReady(false);

      try {
        await createNewScatterplot(controller);
        if (!controller.signal.aborted) {
          scatterplotRef.current = globalScatterplot;
          setScatterplotReady(isZoomReady());
        }
      } catch (error) {
        if (!controller.signal.aborted) {
          console.error('Scatterplot initialization failed:', error);
          globalScatterplot = null;
          setScatterplotReady(false);
        }
      }
    };

    const createNewScatterplot = async (controller: AbortController) => {
      const container = containerRef.current;
      const width = container?.clientWidth || 1000;
      const height = container?.clientHeight || 1000;

      globalScatterplot = new Scatterplot(
        '#sunoverse-container',
        width,
        height
      );

      if (controller.signal.aborted) throw new Error('Aborted');

      await globalScatterplot.load_deeptable({
        source_url: 'https://cdn-o.suno.com/sunoverse/tiles',
      });

      if (controller.signal.aborted) throw new Error('Aborted');

      const rootExtent = globalScatterplot.deeptable.root_tile.extent;

      // Get genre embeddings
      const genreResponse = await fetch(
        'https://cdn-o.suno.com/sunoverse/genre_umap_embs.json'
      );
      const genreData = await genreResponse.json();
      const genreLabels = genreData.map((g: any) => ({
        text: g[0],
        x: g[1],
        y: g[2],
      }));
      setGenreLabels(genreLabels);

      if (controller.signal.aborted) throw new Error('Aborted');

      // Wait for DOM
      await new Promise((resolve) => requestAnimationFrame(resolve));

      const svgElement = document.querySelector(
        '#sunoverse-container #deepscatter-svg'
      );
      if (!svgElement) {
        await new Promise((resolve) => setTimeout(resolve, 100));
      }
      await globalScatterplot.plotAPI({
        encoding: {
          x: { field: 'x' },
          y: { field: 'y' },
          color: {
            field: 'play_count',
            transform: 'log',
            domain: [1, 1000],
            range: 'viridis',
          },
        },
        point_size: 2.0,
        max_points: 1 << 14,
        zoom: { bbox: rootExtent },
        duration: 0,
        labels: {
          name: 'default-labels',
          labels: genreLabels,
          options: LABEL_OPTIONS,
        },
      });

      if (controller.signal.aborted) throw new Error('Aborted');

      const svgEl = document.querySelector(
        '#sunoverse-container #deepscatter-svg'
      ) as SVGElement;
      if (!svgEl) throw new Error('SVG element not found');

      highlightsRef.current = new Highlights(svgEl, globalScatterplot);
      setupCallbacks(globalScatterplot);
    };

    const setupCallbacks = (scatterplot: Scatterplot) => {
      scatterplot.zoom.zoom_callback = () => {
        highlightsRef.current?.updatePoints();

        if (!hoveredPointQidRef.current) return;
        const { x, y } = scatterplot.zoom.qid2screen(
          hoveredPointQidRef.current
        );
        if (!x || !y) return;
        setTooltipPositionRef.current([x + 30, y + 30]);
      };

      scatterplot.label_click = (
        labelData: Record<string, unknown>,
        scatterplot: Scatterplot
      ) => {
        const { x, y } = labelData as any;

        // Zoom to the label location
        scatterplot.zoom.zoom_to(100.0, x, y, 2000);
      };

      scatterplot.zoom.mouseover_callback = async (qid: any) => {
        const { x, y } = scatterplot.zoom.qid2screen(qid);

        if (!qid) {
          hoveredPointQidRef.current = null;
          setTooltipDataRef.current(null);
          highlightsRef.current?.removeTransientCircles(false);
        } else {
          setTooltipPositionRef.current([x + 30, y + 30]);

          // If identical point, no updates
          if (hoveredPointQidRef.current?.[1] === qid?.[1]) return;
          highlightsRef.current?.removeTransientCircles(false);
          hoveredPointQidRef.current = qid;

          const [tix, _] = qid;
          const tile = scatterplot.deeptable.flatTree[tix];
          await tile?.require_columns([
            'play_count',
            'clip_id',
            'image_s3_id',
            'title',
          ]);

          const row = scatterplot.deeptable.getQids([qid])[0];
          const clipId = uint8ArrayToUuid(row.clip_id);

          highlightsRef.current?.addCircle(row, clipId, qid, true, false, {
            radius: 20,
            strokeWidth: 3,
            opacity: 0.8,
          });

          setTooltipDataRef.current({
            title: row.title,
            playCount: row.play_count,
            imageS3Id: row.image_s3_id,
            clipId,
            meta: null,
          });

          // Hydrate with song metadata
          const clipPromise = clipsRef.current.loadClipById(clipId);
          clipPromises.current.set(clipId, clipPromise);
          const clip = await clipPromise;
          clipsRef.current.updateClips([clip]);

          // Do not update data if hovered point changed
          // Could have changed if another callback was called
          if (hoveredPointQidRef.current?.[1] !== qid?.[1]) return;

          setTooltipDataRef.current((data: any) => ({
            ...data,
            meta: {
              display_name: clip?.display_name,
              display_tags: clip?.metadata?.tags,
              user_id: clip?.user_id,
              handle: clip?.handle,
            },
          }));
        }
      };

      scatterplot.click_function = (row, qid) => {
        try {
          zoomToSong(row, qid);
        } catch (error) {
          console.error('Error zooming to song:', error);
        }
      };
    };

    // Set up resize observer
    const resizeObserver = new ResizeObserver(([entry]) => {
      const { width, height } = entry.contentRect;
      if (width > 0 && height > 0)
        scatterplotRef.current?.resize(width, height);
    });
    resizeObserver.observe(containerRef.current);

    initializeScatterplot();

    return () => {
      resizeObserver.disconnect();
      scatterplotRef.current = null;
      highlightsRef.current?.destroy();
      highlightsRef.current = null;
      setScatterplotReady(false);
    };
  }, []);

  // Zoom to song
  const zoomToSong = useCallback(
    async (row: StructRowProxy, qid: Qid, duration = 2000) => {
      if (!scatterplotRef.current) return;
      const scatterplot = scatterplotRef.current;
      const clipId = uint8ArrayToUuid(row.clip_id);

      highlightsRef.current?.removeTransientCircles(true);
      highlightsRef.current?.addCircle(row, clipId, qid, false, true, {
        radius: 25,
        strokeWidth: 3,
        opacity: 1.0,
      });

      scatterplot.zoom.zoom_to(500.0, row.x, row.y, duration);

      // Load song using clipPromises
      // TODO: check if this is cached
      const clip = await clipsRef.current.loadClipById(clipId);
      clipsRef.current.updateClips([clip]);

      if (playbarRef.current.clip?.id === clipId) {
        playbarRef.current.togglePlay();
      } else if (clip) {
        playbarRef.current.playClip(clip);
      } else {
        playbarRef.current.togglePlay(false);
      }
    },
    // TODO: are these all *really* needed? Best to check
    [scatterplotRef, clipsRef, playbarRef]
  );

  const zoomToRandom = useCallback(async () => {
    if (!scatterplotRef.current) return;
    const scatterplot = scatterplotRef.current;

    // Get all loaded tiles that have data
    const loadedTiles = scatterplot.deeptable
      .map((tile) => tile)
      .filter((tile) => tile.record_batch && tile.record_batch.numRows > 0);
    if (loadedTiles.length === 0) {
      console.warn('No loaded tiles with data available');
      return;
    }

    // Pick a random tile and a random row index within that tile
    const randomTile =
      loadedTiles[Math.floor(Math.random() * loadedTiles.length)];
    const randomRowIndex = Math.floor(
      Math.random() * randomTile.record_batch.numRows
    );

    // Create a Qid (tile index, row index pair)
    const qid: Qid = [randomTile.tix, randomRowIndex];
    await randomTile.require_columns(['clip_id', 'image_s3_id', 'x', 'y']);

    // Get the actual row data
    const rows = scatterplot.deeptable.getQids([qid]);
    zoomToSong(rows[0], qid, 5000);
  }, [zoomToSong]);

  const locateMySongs = useCallback(async () => {
    setIsLocatingMyself(true);
    const { data } = await apiClient.GET('/api/sunoverse/embed-my-songs');

    // Get user profile
    if (data) {
      const [x, y] = data[0] as [number, number]; // TODO: figure out what's going on with types here
      scatterplotRef.current?.zoom.zoom_to(500.0, x, y, 2000);

      // TODO: add user profile picture
      highlightsRef.current?.addCircle(
        { x, y, image_url: session.user.avatar_image_url },
        'user-profile',
        null,
        false, // hover transient
        false, // click transient
        {
          radius: 20,
          strokeWidth: 3,
          opacity: 1.0,
        }
      );
    } else {
      console.error('Failed to locate my songs');
    }
    setIsLocatingMyself(false);
  }, [apiClient]);

  const locateGenre = useCallback(
    async (genre: string) => {
      if (!genre) return;

      setIsLocatingGenre(true);
      if (!scatterplotRef.current) return;
      const scatterplot = scatterplotRef.current;

      let genreEmbedding = genreLabels.find((g) => g.text === genre);
      if (!genreEmbedding) {
        // Query server
        const { data } = await apiClient.GET('/api/sunoverse/embed-genres', {
          params: { query: { genres: [genre] } },
        });
        if (!data) {
          // TODO: add toast
          setIsLocatingGenre(false);
          return;
        }

        const coords = (data as any)[0];
        genreEmbedding = { text: genre, x: coords[0], y: coords[1] };
        scatterplot.add_api_label({
          labels: [...genreLabels, genreEmbedding],
          name: `${genre}-label`,
          options: LABEL_OPTIONS,
        });
      }

      console.log('zooming to genre', genreEmbedding);
      scatterplot.zoom.zoom_to(50.0, genreEmbedding.x, genreEmbedding.y, 2000);
      setIsLocatingGenre(false);
    },
    [genreLabels]
  );

  const resetZoom = useCallback(async () => {
    if (!scatterplotRef.current) return;
    const scatterplot = scatterplotRef.current;
    scatterplot.plotAPI({
      zoom: { bbox: scatterplot.deeptable.root_tile.extent },
      duration: 2000,
    });
  }, [scatterplotRef]);

  return (
    <div className='relative h-full w-full overflow-x-hidden'>
      <div
        className='absolute top-0 left-0 h-full w-full'
        id='sunoverse-container'
        ref={containerRef}
      />
      {
        <div
          className={clsx(
            'absolute inset-0 flex h-full w-full flex-col items-center justify-center bg-background-primary transition-opacity duration-1000',
            scatterplotReady ? 'pointer-events-none opacity-0' : 'opacity-100'
          )}
        >
          <p className='mb-2 text-2xl'>Loading...</p>
          <img
            className='rounded-full'
            src={staticAssetUrl('sunoverse/loading.gif')}
            alt='Loading...'
            width={100}
            height={100}
          />
        </div>
      }
      <SongTooltip clipData={tooltipData} position={tooltipPosition} />

      {/* Controls */}
      <div className='absolute top-2 right-2 flex max-w-72 flex-col gap-2 rounded-lg bg-background-primary/10 p-4 backdrop-blur-sm'>
        <p className='mb-2 font-mono text-2xl leading-none font-black uppercase'>
          Sunoverse
        </p>

        <Button onClick={() => resetZoom()}>Reset zoom</Button>
        <Button onClick={locateMySongs}>
          Locate myself
          {isLocatingMyself && <SpinnerSVG className='ml-2 h-4 w-4' />}
        </Button>
        <Button onClick={zoomToRandom}>I&apos;m feeling lucky</Button>
        <Button
          onClick={() => {
            highlightsRef.current?.removeAllCircles();
          }}
        >
          Clear highlights
        </Button>

        <p className='text-sm text-foreground-secondary'>
          Colormap: play count (purple → yellow) (subject to change)
        </p>

        <h2 className='mt-2 text-lg font-bold'>Find genre</h2>
        <div className='flex gap-2'>
          <TextInput
            className='border border-border-primary'
            placeholder='pop, rock, country, etc.'
            value={genreSearchText}
            onChange={(e) => setGenreSearchText(e.target.value)}
            onKeyDown={(e) => e.key === 'Enter' && locateGenre(genreSearchText)}
          />
          <Button
            className='flex items-center justify-center'
            icon={
              isLocatingGenre ? (
                <SpinnerSVG className='size-4' />
              ) : (
                <SearchIcon className='size-4' />
              )
            }
            variant={ButtonVariant.Standard}
            onClick={() => locateGenre(genreSearchText)}
            disabled={isLocatingGenre}
          ></Button>
        </div>
      </div>
    </div>
  );
}