// Copyright (C) 2024 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import {TrackNode} from '../../public/workspace';
import {Trace} from '../../public/trace';
import {PerfettoPlugin} from '../../public/plugin';
import {TrackDescriptor} from '../../public/track';
import {z} from 'zod';
import {
  base64Decode,
  base64Encode,
  utf8Decode,
  utf8Encode,
} from '../../base/string_utils';
import {assertExists} from '../../base/logging';

const PLUGIN_ID = 'dev.perfetto.RestorePinnedTrack';
const SAVED_TRACKS_KEY = `${PLUGIN_ID}#savedPerfettoTracks`;

const RESTORE_COMMAND_ID = `${PLUGIN_ID}#restore`;

/**
 * Fuzzy save and restore of pinned tracks.
 *
 * Tries to persist pinned tracks. Uses full string matching between track name
 * and group name. When no match is found for a saved track, it tries again
 * without numbers.
 */
export default class implements PerfettoPlugin {
  static readonly id = PLUGIN_ID;
  private ctx!: Trace;

  async onTraceLoad(ctx: Trace): Promise<void> {
    this.ctx = ctx;

    ctx.commands.registerCommand({
      id: `${PLUGIN_ID}#save`,
      name: 'Save: Pinned tracks',
      callback: () => {
        this.savedState = {
          ...this.savedState,
          tracks: this.getCurrentPinnedTracks(),
        };
      },
    });
    ctx.commands.registerCommand({
      id: RESTORE_COMMAND_ID,
      name: 'Restore: Pinned tracks',
      callback: () => {
        const tracks = this.savedState?.tracks;
        if (!tracks) {
          alert('No saved tracks. Use the Save command first');
          return;
        }
        this.restoreTracks(tracks);
      },
    });

    ctx.commands.registerCommand({
      id: `${PLUGIN_ID}#saveByName`,
      name: 'Save by name: Pinned tracks',
      callback: async () => {
        const res = await this.ctx.omnibox.prompt(
          'Give a name to the pinned set of tracks',
        );
        if (res) {
          const rawTracksByName = this.savedState?.tracksByName ?? [];
          const tracksByNameMap = new Map(
            rawTracksByName.map((x) => [x.name, x.tracks]),
          );
          tracksByNameMap.set(
            base64Encode(utf8Encode(res)),
            this.getCurrentPinnedTracks(),
          );
          this.savedState = {
            ...this.savedState,
            tracksByName: Array.from(tracksByNameMap.entries()).map(
              ([k, v]) => ({
                name: k,
                tracks: v,
              }),
            ),
          };
        }
      },
    });
    ctx.commands.registerCommand({
      id: `${PLUGIN_ID}#restoreByName`,
      name: 'Restore by name: Pinned tracks',
      callback: async () => {
        const tracksByName = this.savedState?.tracksByName ?? [];
        if (tracksByName.length === 0) {
          alert('No saved tracks. Use the Save by name command first');
          return;
        }
        const res = await this.ctx.omnibox.prompt(
          'Select name of set of pinned tracks to restore',
          tracksByName.map((x) => ({
            key: x.name,
            displayName: utf8Decode(base64Decode(x.name)),
          })),
        );
        if (res) {
          const tracks = assertExists(
            tracksByName.find((x) => x.name === res)?.tracks,
          );
          this.restoreTracks(tracks);
        }
      },
    });
  }

  private restoreTracks(tracks: ReadonlyArray<SavedPinnedTrack>) {
    const localTracks = this.ctx.workspace.flatTracks.map((track) => ({
      savedTrack: this.toSavedTrack(track),
      track: track,
    }));
    tracks.forEach((trackToRestore) => {
      const foundTrack = this.findMatchingTrack(localTracks, trackToRestore);
      if (foundTrack) {
        foundTrack.pin();
      } else {
        console.warn(
          '[RestorePinnedTracks] No track found that matches',
          trackToRestore,
        );
      }
    });
  }

  private getCurrentPinnedTracks() {
    return this.ctx.workspace.pinnedTracks.map((track) =>
      this.toSavedTrack(track),
    );
  }

  private findMatchingTrack(
    localTracks: Array<LocalTrack>,
    savedTrack: SavedPinnedTrack,
  ): TrackNode | undefined {
    let mostSimilarTrack: LocalTrack | undefined = undefined;
    let mostSimilarTrackDifferenceScore: number = 0;

    for (let i = 0; i < localTracks.length; i++) {
      const localTrack = localTracks[i];
      const differenceScore = this.calculateSimilarityScore(
        localTrack.savedTrack,
        savedTrack,
      );

      // Return immediately if we found the exact match
      if (differenceScore === Number.MAX_SAFE_INTEGER) {
        return localTrack.track;
      }

      // Ignore too different objects
      if (differenceScore === 0) {
        continue;
      }

      if (differenceScore > mostSimilarTrackDifferenceScore) {
        mostSimilarTrackDifferenceScore = differenceScore;
        mostSimilarTrack = localTrack;
      }
    }

    return mostSimilarTrack?.track || undefined;
  }

  /**
   * Returns the similarity score where 0 means the objects are completely
   * different, and the higher the number, the smaller the difference is.
   * Returns Number.MAX_SAFE_INTEGER if the objects are completely equal.
   * We attempt a fuzzy match based on the similarity score.
   * For example, one of the ways we do this is we remove the numbers
   * from the title to potentially pin a "similar" track from a different trace.
   * Removing numbers allows flexibility; for instance, with multiple 'sysui'
   * processes (e.g. track group name: "com.android.systemui 123") without
   * this approach, any could be mistakenly pinned. The goal is to restore
   * specific tracks within the same trace, ensuring that a previously pinned
   * track is pinned again.
   * If the specific process with that PID is unavailable, pinning any
   * other process matching the package name is attempted.
   * @param track1 first saved track to compare
   * @param track2 second saved track to compare
   * @private
   */
  private calculateSimilarityScore(
    track1: SavedPinnedTrack,
    track2: SavedPinnedTrack,
  ): number {
    // Return immediately when objects are equal
    if (
      track1.trackName === track2.trackName &&
      track1.groupName === track2.groupName &&
      track1.pluginId === track2.pluginId &&
      track1.kind === track2.kind &&
      track1.isMainThread === track2.isMainThread
    ) {
      return Number.MAX_SAFE_INTEGER;
    }

    let similarityScore = 0;
    if (track1.trackName === track2.trackName) {
      similarityScore += 100;
    } else if (
      this.removeNumbers(track1.trackName) ===
      this.removeNumbers(track2.trackName)
    ) {
      similarityScore += 50;
    }

    if (track1.groupName === track2.groupName) {
      similarityScore += 90;
    } else if (
      this.removeNumbers(track1.groupName) ===
      this.removeNumbers(track2.groupName)
    ) {
      similarityScore += 45;
    }

    // Do not consider other parameters if there is no match in name/group
    if (similarityScore === 0) return similarityScore;

    if (track1.pluginId === track2.pluginId) {
      similarityScore += 30;
    }

    if (track1.kind === track2.kind) {
      similarityScore += 20;
    }

    if (track1.isMainThread === track2.isMainThread) {
      similarityScore += 10;
    }

    return similarityScore;
  }

  private removeNumbers(inputString?: string): string | undefined {
    return inputString?.replace(/\d+/g, '');
  }

  private toSavedTrack(track: TrackNode): SavedPinnedTrack {
    let trackDescriptor: TrackDescriptor | undefined = undefined;
    if (track.uri != undefined) {
      trackDescriptor = this.ctx.tracks.getTrack(track.uri);
    }

    return {
      groupName: groupName(track),
      trackName: track.title,
      pluginId: trackDescriptor?.pluginId,
      kind: trackDescriptor?.tags?.kind,
      isMainThread: trackDescriptor?.chips?.includes('main thread') || false,
    };
  }

  private get savedState(): SavedState | undefined {
    const savedStateString = window.localStorage.getItem(SAVED_TRACKS_KEY);
    if (!savedStateString) {
      return undefined;
    }
    const savedState = SAVED_STATE_SCHEMA.safeParse(
      JSON.parse(savedStateString),
    );
    if (!savedState.success) {
      return undefined;
    }
    return savedState.data;
  }

  private set savedState(state: SavedState) {
    window.localStorage.setItem(SAVED_TRACKS_KEY, JSON.stringify(state));
  }
}

// Return the displayname of the containing group
// If the track is a child of a workspace, return undefined...
function groupName(track: TrackNode): string | undefined {
  const parent = track.parent;
  if (parent instanceof TrackNode) {
    return parent.title;
  }
  return undefined;
}

const SAVED_PINNED_TRACK_SCHEMA = z
  .object({
    // Optional: group name for the track. Usually matches with process name.
    groupName: z.string().optional(),
    // Track name to restore.
    trackName: z.string(),
    // Plugin used to create this track
    pluginId: z.string().optional(),
    // Kind of the track
    kind: z.string().optional(),
    // If it's a thread track, it should be true in case it's a main thread track
    isMainThread: z.boolean(),
  })
  .readonly();

type SavedPinnedTrack = z.infer<typeof SAVED_PINNED_TRACK_SCHEMA>;

const SAVED_STATE_SCHEMA = z
  .object({
    tracks: z.array(SAVED_PINNED_TRACK_SCHEMA).optional().readonly(),
    tracksByName: z
      .array(
        z
          .object({
            name: z.string(),
            tracks: z.array(SAVED_PINNED_TRACK_SCHEMA).readonly(),
          })
          .readonly(),
      )
      .optional()
      .readonly(),
  })
  .readonly();

type SavedState = z.infer<typeof SAVED_STATE_SCHEMA>;

interface LocalTrack {
  readonly savedTrack: SavedPinnedTrack;
  readonly track: TrackNode;
}