import { deduplicateMarkersArray } from '@suno/studiokit/projectState/warpMarkersRegistry';

import { StudioClip, StudioProjectState, StudioTakeLane } from './types';

type WarpMarkers = StudioClip['warp']['markers'];

type BackwardCompatibleClip = Omit<StudioClip, 'warp'> & {
  warp: Omit<StudioClip['warp'], 'markers'> & {
    markers?: WarpMarkers;
    markersHash?: string;
  };
};
export interface SerializedProjectState
  extends Omit<StudioProjectState, 'tracks'> {
  // Registry of deduplicated markers indexed by hash
  markersRegistry: Record<string, WarpMarkers>;

  tracks: Array<
    Omit<StudioProjectState['tracks'][0], 'clips' | 'takeLanes'> & {
      clips: Array<BackwardCompatibleClip>;
      takeLanes: Array<
        Omit<StudioTakeLane, 'clips'> & {
          clips: Array<BackwardCompatibleClip>;
        }
      >;
    }
  >;
}

const serializeClip = (
  mapping: Map<WarpMarkers, WarpMarkers>,
  hashMapping: Map<WarpMarkers, string>,
  clip: StudioClip
) => {
  if (!clip.warp || !clip.warp.markers) {
    return clip;
  }

  const canonicalMarkers = mapping.get(clip.warp.markers);
  if (!canonicalMarkers) {
    return clip;
  }

  const hash = hashMapping.get(canonicalMarkers);
  if (!hash) {
    return clip;
  }

  const { markers, ...warpWithoutMarkers } = clip.warp;
  return {
    ...clip,
    warp: {
      ...warpWithoutMarkers,
      markersHash: hash,
    },
  };
};

/**
 * Serialize a project state with deduplicated warp markers.
 * Creates a markersRegistry containing unique markers indexed by hash,
 * and replaces clip.warp.markers with clip.warp.markersHash references.
 */
export function serializeProjectState(
  state: StudioProjectState
): SerializedProjectState {
  // Collect all unique markers from all clips
  const allWarpMarkers: WarpMarkers[] = [];

  state.tracks.forEach((track) => {
    track.clips.forEach((clip) => {
      if (clip.warp && clip.warp.markers) {
        allWarpMarkers.push(clip.warp.markers);
      }
    });
    track.takeLanes.forEach((takeLane) => {
      takeLane.clips.forEach((clip) => {
        if (clip.warp && clip.warp.markers) {
          allWarpMarkers.push(clip.warp.markers);
        }
      });
    });
  });

  // Deduplicate markers and get hash mapping
  const { deduplicated, mapping, hashMapping } =
    deduplicateMarkersArray(allWarpMarkers);

  // Build the registry
  const markersRegistry: Record<string, WarpMarkers> = {};
  deduplicated.forEach((markers) => {
    const hash = hashMapping.get(markers)!;
    markersRegistry[hash] = markers;
  });

  // Transform the state to use hash references
  const serializedState: SerializedProjectState = {
    ...state,
    markersRegistry,
    tracks: state.tracks.map((track) => ({
      ...track,
      takeLanes: track.takeLanes.map((takeLane) => ({
        ...takeLane,
        clips: takeLane.clips.map((clip) =>
          serializeClip(mapping, hashMapping, clip)
        ),
      })),
      clips: track.clips.map((clip) =>
        serializeClip(mapping, hashMapping, clip)
      ),
    })),
  };

  return serializedState;
}

const deserializeClip = (
  markersRegistry: Record<string, WarpMarkers>,
  clip: BackwardCompatibleClip
) => {
  const { markersHash, markers, ...warpWithoutHashAndMarkers } = clip.warp;

  // If clip has markersHash, look it up in registry
  if (markersHash && markersRegistry[markersHash]) {
    return {
      ...clip,
      warp: {
        ...warpWithoutHashAndMarkers,
        markers: markersRegistry[markersHash],
      },
    } as StudioClip;
  }

  // Fallback to direct markers (backward compatibility or missing hash)
  return {
    ...clip,
    warp: {
      ...warpWithoutHashAndMarkers,
      markers: markers || {},
    },
  } as StudioClip;
};

/**
 * Deserialize a project state, converting markersHash references back to markers objects.
 * This function handles both the new format (with markersRegistry) and the old format
 * (with direct markers) for backward compatibility.
 */
export function deserializeProjectState(
  serializedState: SerializedProjectState
): StudioProjectState {
  const { markersRegistry, ...restState } = serializedState;

  // Transform clips back to use direct markers references
  const deserializedState: StudioProjectState = {
    ...restState,
    tracks: serializedState.tracks.map((track) => ({
      ...track,
      takeLanes: (track.takeLanes || []).map((takeLane) => ({
        ...takeLane,
        clips: takeLane.clips.map((clip) =>
          deserializeClip(markersRegistry, clip)
        ),
      })),
      clips: track.clips.map((clip) => deserializeClip(markersRegistry, clip)),
    })),
  };

  return deserializedState;
}

export function isSerializedProjectState(
  state: SerializedProjectState | StudioProjectState
): state is SerializedProjectState {
  return 'markersRegistry' in state;
}

export function maybeDeserializeProjectState(
  serializedState: SerializedProjectState | StudioProjectState
): StudioProjectState {
  if (isSerializedProjectState(serializedState)) {
    return deserializeProjectState(serializedState);
  }
  return serializedState;
}

export function maybeDeserializeProjectStateInPlace(
  serializedState: SerializedProjectState | StudioProjectState
): StudioProjectState {
  if (isSerializedProjectState(serializedState)) {
    const deserialized = deserializeProjectState(serializedState);
    Object.keys(serializedState).forEach(
      (key) => delete (serializedState as any)[key]
    );
    Object.keys(deserialized).forEach(
      (key) =>
        ((serializedState as any)[key] =
          deserialized[key as keyof StudioProjectState])
    );
  }
  return serializedState as StudioProjectState;
}
