// This file is part of the 64k demo project.
// It implements the spectool for analyzing audio into spectrograms.
// Provides both 'analyze' and 'play' modes for spectral data.

#include "audio/audio.h"
#include "audio/dct.h"
#include "audio/gen.h"
#include "audio/synth.h"
#include "audio/window.h"
#include "platform.h"
#include <stdio.h>
#include <string.h>

#include "miniaudio.h"

#include <math.h>
#include <stdlib.h>
#include <time.h>
#include <vector>

// Simple .spec file format:
// char[4] magic = "SPEC"
// int32_t version = 1
// int32_t dct_size
// int32_t num_frames
// float[num_frames * dct_size] data
struct SpecHeader {
  char magic[4];
  int32_t version;
  int32_t dct_size;
  int32_t num_frames;
};

int analyze_audio(const char* in_path, const char* out_path) {
  printf("Analyzing %s -> %s\n", in_path, out_path);

  ma_decoder_config config = ma_decoder_config_init(ma_format_f32, 1, 32000);
  ma_decoder decoder;
  if (ma_decoder_init_file(in_path, &config, &decoder) != MA_SUCCESS) {
    printf("Error: Failed to open or decode audio file: %s\n", in_path);
    return 1;
  }

  std::vector<float> spec_data;
  float pcm_chunk[DCT_SIZE];
  float window[WINDOW_SIZE];
  hamming_window_512(window);

  ma_uint64 frames_read;
  while (ma_decoder_read_pcm_frames(&decoder, pcm_chunk, DCT_SIZE,
                                    &frames_read) == MA_SUCCESS &&
         frames_read > 0) {
    if (frames_read < DCT_SIZE) {
      // Zero-pad the last chunk if it's smaller
      memset(pcm_chunk + frames_read, 0,
             (DCT_SIZE - frames_read) * sizeof(float));
    }

    // Apply window
    for (int i = 0; i < DCT_SIZE; ++i) {
      pcm_chunk[i] *= window[i];
    }

    // Apply FDCT
    float dct_chunk[DCT_SIZE];
    fdct_512(pcm_chunk, dct_chunk);

    // Add to spectrogram data
    spec_data.insert(spec_data.end(), dct_chunk, dct_chunk + DCT_SIZE);
  }

  ma_decoder_uninit(&decoder);

  // Trim trailing zero frames
  int last_frame = spec_data.size() / DCT_SIZE;
  while (last_frame > 0) {
    bool all_zeros = true;
    for (int i = 0; i < DCT_SIZE; ++i) {
      if (spec_data[(last_frame - 1) * DCT_SIZE + i] != 0.0f) {
        all_zeros = false;
        break;
      }
    }
    if (all_zeros) {
      --last_frame;
    } else {
      break;
    }
  }
  spec_data.resize(last_frame * DCT_SIZE);

  // Write to .spec file
  FILE* f_out = fopen(out_path, "wb");
  if (!f_out) {
    printf("Error: Failed to open output file: %s\n", out_path);
    return 1;
  }

  SpecHeader header;
  memcpy(header.magic, "SPEC", 4);
  header.version = 1;
  header.dct_size = DCT_SIZE;
  header.num_frames = spec_data.size() / DCT_SIZE;

  fwrite(&header, sizeof(SpecHeader), 1, f_out);
  fwrite(spec_data.data(), sizeof(float), spec_data.size(), f_out);
  fclose(f_out);

  printf("Analysis complete. Wrote %d spectral frames.\n", header.num_frames);
  return 0;
}

int play_spec(const char* in_path) {
  printf("Playing %s\n", in_path);

  FILE* f_in = fopen(in_path, "rb");
  if (!f_in) {
    printf("Error: Failed to open input file: %s\n", in_path);
    return 1;
  }

  SpecHeader header;
  if (fread(&header, sizeof(SpecHeader), 1, f_in) != 1 ||
      strncmp(header.magic, "SPEC", 4) != 0) {
    printf("Error: Invalid spectrogram file format.\n");
    fclose(f_in);
    return 1;
  }

  std::vector<float> spec_data(header.num_frames * header.dct_size);
  fread(spec_data.data(), sizeof(float), spec_data.size(), f_in);
  fclose(f_in);

  platform_init_window(false);
  audio_init();
  audio_start();
  Spectrogram spec;
  spec.spectral_data_a = spec_data.data();
  spec.spectral_data_b =
      spec_data.data(); // Point both to the same buffer for playback
  spec.num_frames = header.num_frames;

  int spec_id = synth_register_spectrogram(&spec);
  synth_trigger_voice(spec_id, 0.7f, 0.0f);

  printf("Playing... Press Ctrl+C to exit.\n");
  while (synth_get_active_voice_count() > 0 && !platform_should_close()) {
    platform_poll();
  }

  audio_shutdown();
  platform_shutdown();

  return 0;
}

int test_gen(const char* out_path) {
  printf("Generating test spectrogram -> %s\n", out_path);

  std::vector<float> track_data;
  int track_frames = 0;

  // Generate a simple C Major scale
  float freqs[] = {261.63f, 293.66f, 329.63f, 349.23f,
                   392.00f, 440.00f, 493.88f, 523.25f};

  srand(time(NULL));

  for (int i = 0; i < 8; ++i) {
    NoteParams params;
    params.base_freq = freqs[i];
    params.duration_sec = 0.5f;
    params.amplitude = 0.5f;
    params.attack_sec = 0.05f;
    params.decay_sec = 0.1f;
    params.vibrato_rate = 5.0f;
    params.vibrato_depth = 2.0f;
    params.num_harmonics = 5;
    params.harmonic_decay = 0.5f;
    params.pitch_randomness = 1.0f;
    params.amp_randomness = 0.05f;

    int note_frames = 0;
    std::vector<float> note_data =
        generate_note_spectrogram(params, &note_frames);

    // Paste at 0.4s intervals (overlap)
    int offset = (int)(i * 0.4f * 32000.0f / DCT_SIZE);
    paste_spectrogram(track_data, &track_frames, note_data, note_frames,
                      offset);
  }

  // Write to file (Duplicate logic, but fine for now)
  FILE* f_out = fopen(out_path, "wb");
  if (!f_out) {
    printf("Error: Failed to open output file: %s\n", out_path);
    return 1;
  }

  SpecHeader header;
  memcpy(header.magic, "SPEC", 4);
  header.version = 1;
  header.dct_size = DCT_SIZE;
  header.num_frames = track_frames;

  fwrite(&header, sizeof(SpecHeader), 1, f_out);
  fwrite(track_data.data(), sizeof(float), track_data.size(), f_out);
  fclose(f_out);

  printf("Generated %d frames.\n", track_frames);
  return 0;
}

void print_usage() {
  printf("Usage: spectool <command> <input> [output]\n");
  printf("Commands:\n");
  printf(
      "  analyze <input.wav|.mp3> <output.spec>   Analyze an audio file and "
      "save as a spectrogram.\n");
  printf(
      "  play    <input.spec>                     Play a spectrogram file.\n");
  printf(
      "  test_gen <output.spec>                   Generate a test "
      "spectrogram.\n");
}

int main(int argc, char** argv) {
  if (argc < 2) {
    print_usage();
    return 1;
  }

  const char* command = argv[1];

  if (strcmp(command, "analyze") == 0) {
    if (argc < 4) {
      printf("Error: 'analyze' command requires input and output files.\n");
      print_usage();
      return 1;
    }
    return analyze_audio(argv[2], argv[3]);
  } else if (strcmp(command, "play") == 0) {
    if (argc < 3) {
      printf("Error: 'play' command requires an input file.\n");
      print_usage();
      return 1;
    }
    return play_spec(argv[2]);
  } else if (strcmp(command, "test_gen") == 0) {
    if (argc < 3) {
      printf("Error: 'test_gen' command requires an output file.\n");
      print_usage();
      return 1;
    }
    return test_gen(argv[2]);
  } else {
    printf("Error: Unknown command '%s'\n", command);
    print_usage();
    return 1;
  }

  return 0;
}