import { SunoClient, SunoLyricsVariant, SunoTaskData, SunoTrack } from './client';
import * as schemas from './schemas';

const ASYNC_POLLING_PATHS: Record<string, string> = {
  'suno_generate_music': '/generate/record-info?taskId=',
  'suno_generate_lyrics': '/lyrics/record-info?taskId=',
  'suno_generate_sounds': '/generate/record-info?taskId=',
  'suno_extend_music': '/generate/record-info?taskId=',
  'suno_replace_section': '/generate/record-info?taskId=',
  'suno_upload_and_cover': '/generate/record-info?taskId=',
  'suno_upload_and_extend': '/generate/record-info?taskId=',
  'suno_add_vocals': '/generate/record-info?taskId=',
  'suno_add_instrumental': '/generate/record-info?taskId=',
  'suno_generate_mashup': '/generate/record-info?taskId=',
  'suno_separate_vocals': '/vocal-removal/record-info?taskId=',
  'suno_separate_stems': '/vocal-removal/record-info?taskId=',
  'suno_create_video': '/generate/video-info?taskId=',
  'suno_create_cover': '/generate/cover-info?taskId=',
  'suno_convert_to_wav': '/generate/wav-info?taskId=',
};

function copyDefined<T extends Record<string, unknown>>(
  source: T,
  target: Record<string, unknown>,
  keys: (keyof T)[]
): void {
  for (const key of keys) {
    const value = source[key];
    if (value !== undefined) {
      target[key as string] = value;
    }
  }
}

function isLyricsData(data: SunoTaskData): boolean {
  return data.type === 'LYRICS';
}

function formatLyricsResult(data: SunoTaskData): string {
  const status = data.status;
  if (status !== 'SUCCESS') {
    return JSON.stringify(
      {
        taskId: data.taskId,
        status,
        errorCode: data.errorCode,
        errorMessage: data.errorMessage,
      },
      null,
      2
    );
  }

  const variants = (data.response?.data as SunoLyricsVariant[] | undefined) ?? [];
  return JSON.stringify(
    {
      taskId: data.taskId,
      status,
      type: data.type,
      variants: variants.map((v, index) => ({
        index: index + 1,
        title: v.title,
        text: v.text,
        status: v.status,
        errorMessage: v.errorMessage,
      })),
    },
    null,
    2
  );
}

function formatTaskResult(data: SunoTaskData): string {
  if (isLyricsData(data)) {
    return formatLyricsResult(data);
  }

  const status = data.status;
  if (status !== 'SUCCESS') {
    return JSON.stringify(
      {
        taskId: data.taskId,
        status,
        errorCode: data.errorCode,
        errorMessage: data.errorMessage,
      },
      null,
      2
    );
  }

  const response = data.response;
  const tracks = (response?.data ?? response?.sunoData ?? []) as SunoTrack[];
  return JSON.stringify(
    {
      taskId: data.taskId,
      status,
      operationType: data.operationType,
      tracks: tracks.map((t) => ({
        id: t.id,
        title: t.title,
        tags: t.tags,
        duration: t.duration,
        prompt: t.prompt,
        audio_url: t.audio_url ?? t.audioUrl,
        stream_audio_url: t.stream_audio_url ?? t.streamAudioUrl,
        image_url: t.image_url ?? t.imageUrl,
      })),
    },
    null,
    2
  );
}

async function runAsyncTask(
  client: SunoClient,
  taskName: string,
  createPath: string,
  payload: Record<string, unknown>
): Promise<string> {
  const createResult = await client.post<{ taskId: string }>(createPath, payload);
  const taskId = createResult.data.taskId;

  const statusPathTemplate = ASYNC_POLLING_PATHS[taskName];
  if (!statusPathTemplate) {
    return JSON.stringify({ taskId, status: 'PENDING', note: 'Use suno_get_task_status to poll.' }, null, 2);
  }

  const statusData = await client.pollTaskStatus(taskId, {
    statusPath: `${statusPathTemplate}${encodeURIComponent(taskId)}`,
  });

  return formatTaskResult(statusData);
}

export class ToolHandlers {
  constructor(private readonly client: SunoClient) {}

  async handle(name: string, args: unknown): Promise<{ content: { type: string; text: string }[] }> {
    let resultText: string;

    switch (name) {
      case 'suno_generate_music': {
        const input = schemas.GenerateMusicSchema.parse(args);
        const payload: Record<string, unknown> = {
          customMode: input.customMode,
          instrumental: input.instrumental,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, [
          'prompt',
          'style',
          'title',
          'negativeTags',
          'vocalGender',
          'styleWeight',
          'weirdnessConstraint',
          'audioWeight',
          'personaId',
          'personaModel',
        ]);
        resultText = await runAsyncTask(this.client, name, '/generate', payload);
        break;
      }

      case 'suno_generate_lyrics': {
        const input = schemas.GenerateLyricsSchema.parse(args);
        resultText = await runAsyncTask(this.client, name, '/lyrics', {
          prompt: input.prompt,
          callBackUrl: input.callBackUrl ?? '',
        });
        break;
      }

      case 'suno_generate_sounds': {
        const input = schemas.GenerateSoundsSchema.parse(args);
        const payload: Record<string, unknown> = {
          prompt: input.prompt,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, ['soundLoop', 'soundTempo', 'soundKey', 'grabLyrics']);
        resultText = await runAsyncTask(this.client, name, '/generate/sounds', payload);
        break;
      }

      case 'suno_extend_music': {
        const input = schemas.ExtendMusicSchema.parse(args);
        const payload: Record<string, unknown> = {
          audioId: input.audioId,
          defaultParamFlag: input.defaultParamFlag,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, [
          'prompt',
          'style',
          'title',
          'continueAt',
          'negativeTags',
          'vocalGender',
          'styleWeight',
          'weirdnessConstraint',
          'audioWeight',
          'personaId',
          'personaModel',
        ]);
        resultText = await runAsyncTask(this.client, name, '/generate/extend', payload);
        break;
      }

      case 'suno_replace_section': {
        const input = schemas.ReplaceSectionSchema.parse(args);
        const payload: Record<string, unknown> = {
          audioId: input.audioId,
          prompt: input.prompt,
          start: input.start,
          end: input.end,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, ['replaceAudioUrl', 'customMode', 'style', 'title', 'instrumental']);
        resultText = await runAsyncTask(this.client, name, '/generate/replace-section', payload);
        break;
      }

      case 'suno_add_vocals': {
        const input = schemas.AddVocalsSchema.parse(args);
        const payload: Record<string, unknown> = {
          audioId: input.audioId,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, ['prompt', 'customMode', 'style', 'title']);
        resultText = await runAsyncTask(this.client, name, '/generate/add-vocals', payload);
        break;
      }

      case 'suno_add_instrumental': {
        const input = schemas.AddInstrumentalSchema.parse(args);
        const payload: Record<string, unknown> = {
          audioId: input.audioId,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, ['prompt', 'customMode', 'style', 'title']);
        resultText = await runAsyncTask(this.client, name, '/generate/add-instrumental', payload);
        break;
      }

      case 'suno_upload_and_cover': {
        const input = schemas.UploadAndCoverSchema.parse(args);
        const payload: Record<string, unknown> = {
          uploadUrl: input.uploadUrl,
          customMode: input.customMode,
          instrumental: input.instrumental,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, [
          'prompt',
          'style',
          'title',
          'negativeTags',
          'vocalGender',
          'styleWeight',
          'weirdnessConstraint',
          'audioWeight',
          'personaId',
          'personaModel',
        ]);
        resultText = await runAsyncTask(this.client, name, '/generate/upload-cover', payload);
        break;
      }

      case 'suno_upload_and_extend': {
        const input = schemas.UploadAndExtendSchema.parse(args);
        const payload: Record<string, unknown> = {
          uploadUrl: input.uploadUrl,
          defaultParamFlag: input.defaultParamFlag,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, ['prompt', 'style', 'title', 'continueAt']);
        resultText = await runAsyncTask(this.client, name, '/generate/upload-extend', payload);
        break;
      }

      case 'suno_generate_mashup': {
        const input = schemas.GenerateMashupSchema.parse(args);
        const payload: Record<string, unknown> = {
          uploadUrlList: input.uploadUrlList,
          customMode: input.customMode,
          prompt: input.prompt,
          model: input.model,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, [
          'instrumental',
          'style',
          'title',
          'negativeTags',
          'vocalGender',
          'styleWeight',
          'weirdnessConstraint',
          'audioWeight',
        ]);
        resultText = await runAsyncTask(this.client, name, '/generate/mashup', payload);
        break;
      }

      case 'suno_separate_vocals': {
        const input = schemas.SeparateVocalsSchema.parse(args);
        resultText = await runAsyncTask(this.client, name, '/vocal-removal/generate', {
          taskId: input.taskId,
          audioId: input.audioId,
          callBackUrl: input.callBackUrl ?? '',
        });
        break;
      }

      case 'suno_separate_stems': {
        const input = schemas.SeparateStemsSchema.parse(args);
        resultText = await runAsyncTask(this.client, name, '/vocal-removal/generate', {
          taskId: input.taskId,
          audioId: input.audioId,
          type: 'split_stem',
          callBackUrl: input.callBackUrl ?? '',
        });
        break;
      }

      case 'suno_create_video': {
        const input = schemas.CreateVideoSchema.parse(args);
        const payload: Record<string, unknown> = {
          taskId: input.taskId,
          audioId: input.audioId,
          callBackUrl: input.callBackUrl ?? '',
        };
        copyDefined(input, payload, ['author', 'domainName']);
        resultText = await runAsyncTask(this.client, name, '/generate/video', payload);
        break;
      }

      case 'suno_create_cover': {
        const input = schemas.CreateCoverSchema.parse(args);
        resultText = await runAsyncTask(this.client, name, '/generate/cover', {
          taskId: input.taskId,
          audioId: input.audioId,
          callBackUrl: input.callBackUrl ?? '',
        });
        break;
      }

      case 'suno_convert_to_wav': {
        const input = schemas.ConvertToWavSchema.parse(args);
        resultText = await runAsyncTask(this.client, name, '/generate/wav', {
          taskId: input.taskId,
          audioId: input.audioId,
          callBackUrl: input.callBackUrl ?? '',
        });
        break;
      }

      case 'suno_boost_style': {
        const input = schemas.BoostStyleSchema.parse(args);
        const result = await this.client.post<{ result: string }>('/style/generate', {
          content: input.content,
        });
        resultText = result.data.result;
        break;
      }

      case 'suno_generate_persona': {
        const input = schemas.GeneratePersonaSchema.parse(args);
        const payload: Record<string, unknown> = { taskId: input.taskId };
        if (input.audioId !== undefined) payload.audioId = input.audioId;
        const result = await this.client.post<{ personaId?: string }>('/generate/persona', payload);
        resultText = JSON.stringify(result.data, null, 2);
        break;
      }

      case 'suno_get_task_status': {
        const input = schemas.GetTaskStatusSchema.parse(args);
        const data = await this.client.get<SunoTaskData>(
          `/generate/record-info?taskId=${encodeURIComponent(input.taskId)}`
        );
        resultText = formatTaskResult(data.data);
        break;
      }

      case 'suno_get_timestamped_lyrics': {
        const input = schemas.GetTimestampedLyricsSchema.parse(args);
        const result = await this.client.get<unknown>(
          `/generate/timestamped-lyrics?taskId=${encodeURIComponent(input.taskId)}&audioId=${encodeURIComponent(
            input.audioId
          )}`
        );
        resultText = JSON.stringify(result.data, null, 2);
        break;
      }

      case 'suno_get_lyrics_details': {
        const input = schemas.GetLyricsDetailsSchema.parse(args);
        const result = await this.client.get<SunoTaskData>(
          `/lyrics/record-info?taskId=${encodeURIComponent(input.taskId)}`
        );
        resultText = formatLyricsResult(result.data);
        break;
      }

      case 'suno_get_wav_details': {
        const input = schemas.GetWavDetailsSchema.parse(args);
        const result = await this.client.get<unknown>(
          `/generate/wav-info?taskId=${encodeURIComponent(input.taskId)}`
        );
        resultText = JSON.stringify(result.data, null, 2);
        break;
      }

      case 'suno_get_vocal_separation_details': {
        const input = schemas.GetVocalSeparationDetailsSchema.parse(args);
        const result = await this.client.get<unknown>(
          `/vocal-removal/record-info?taskId=${encodeURIComponent(input.taskId)}`
        );
        resultText = JSON.stringify(result.data, null, 2);
        break;
      }

      case 'suno_get_video_details': {
        const input = schemas.GetVideoDetailsSchema.parse(args);
        const result = await this.client.get<unknown>(
          `/generate/video-info?taskId=${encodeURIComponent(input.taskId)}`
        );
        resultText = JSON.stringify(result.data, null, 2);
        break;
      }

      case 'suno_get_credits': {
        schemas.GetCreditsSchema.parse(args);
        const result = await this.client.get<number>('/generate/credit');
        resultText = JSON.stringify({ credits: result.data }, null, 2);
        break;
      }

      default:
        throw new Error(`Unknown tool: ${name}`);
    }

    return {
      content: [{ type: 'text', text: resultText }],
    };
  }
}