import * as fs from 'fs';
import * as https from 'https';
import chalk from 'chalk';
import * as dotenv from 'dotenv';
dotenv.config();

let songCount = 1;

const PROGRESS_FILE = './progress.json';

type ProgressState = {
  songCount: number;
  currentMentalStateIndex: number;
  currentGenreIndex: number;
  processedSongs: Set<string>;
  completedGenres: string[];
};

function loadProgress(): ProgressState | null {
  if (fs.existsSync(PROGRESS_FILE)) {
    try {
      const data = JSON.parse(fs.readFileSync(PROGRESS_FILE, 'utf-8'));
      console.log(chalk.blue('Found previous progress, resuming...'));
      console.log(chalk.blue(`Previously downloaded: ${data.songCount} songs`));
      return {
        ...data,
        processedSongs: new Set(data.processedSongs || [])
      };
    } catch (error) {
      console.error(chalk.red('Error loading progress file, starting fresh'), error);
      return null;
    }
  }
  return null;
}

function saveProgress(state: ProgressState) {
  const dataToSave = {
    ...state,
    processedSongs: Array.from(state.processedSongs)
  };
  fs.writeFileSync(PROGRESS_FILE, JSON.stringify(dataToSave, null, 2));
}

function clearProgress() {
  if (fs.existsSync(PROGRESS_FILE)) {
    fs.unlinkSync(PROGRESS_FILE);
    console.log(chalk.green('Progress file cleared - all downloads complete!'));
  }
}

const AUTHTOKEN = process.env.AUTHTOKEN;

if (!AUTHTOKEN) {
  console.error('Please enter all the required information');
  process.exit(1);
}

type MentalState = "focus" | "relax" | "sleep" | "meditate";

type MentalStateArray = MentalState[];

let mentalStates: MentalStateArray = ["focus", "relax", "sleep", "meditate"];

type GenreStructure = {
  base: string[];
  nature: string[];
};

type GenresMap = Record<MentalState, GenreStructure>;

let genres: GenresMap = {
  "focus": {
    base: [
      "Acoustic",
      "Atmospheric",
      "Cinematic",
      "Classical",
      "Drone",
      "Electronic",
      "Grooves",
      "Lofi",
      "Piano",
      "Post Rock"
    ],
    nature: [
      "Beach",
      "Chimes & Bowls",
      "Forest",
      "Nightsounds",
      "Rain",
      "Rainforest",
      "River",
      "Thunder",
      "Underwater",
      "Wind"
    ]
  },
  "relax": {
    base: [
      "Atmospheric",
      "Electronic"
    ],
    nature: [
      "Beach",
      "Chimes & Bowls",
      "Forest",
      "Nightsounds",
      "Rain",
      "Rainforest",
      "River",
      "Thunder",
      "Underwater",
      "Wind"
    ]
  },
  "sleep": {
    base: [
      "Atmospheric"
    ],
    nature: [
      "Beach",
      "Forest",
      "Nightsounds",
      "Rain",
      "Rainforest",
      "River",
      "Thunder",
      "Underwater",
      "Wind"
    ]
  },
  "meditate": {
    base: [
      "Atmospheric",
      "Electronic"
    ],
    nature: [
      "Beach",
      "Chimes & Bowls",
      "Forest",
      "Nightsounds",
      "Rain",
      "Rainforest",
      "River",
      "Thunder",
      "Underwater",
      "Wind"
    ]
  }
};

// Commenting out moods since it's not used in the code
/*
type MoodsMap = Record<MentalState, string[]>;

let moods: MoodsMap = {
  "focus": [
    "Brooding",
    "Calm",
    "Chill",
    "Dark",
    "Downtempo",
    "Dreamlike",
    "Driving",
    "Energizing",
    "Epic",
    "Floating",
    "Heavy",
    "Hopeful",
    "Inspiring",
    "Meditative",
    "Mysterious",
    "Ominous",
    "Optimistic",
    "Playful",
    "Ponderous",
    "Serene",
    "Strong",
    "Upbeat",
    "Uplifting"
  ],
  "relax": [
    "Brooding",
    "Calm",
    "Chill",
    "Dark",
    "Downtempo",
    "Dreamlike",
    "Driving",
    "Energizing",
    "Epic",
    "Floating",
    "Hopeful",
    "Inspiring",
    "Meditative",
    "Mysterious",
    "Optimistic",
    "Playful",
    "Ponderous",
    "Serene",
    "Strong",
    "Upbeat",
    "Uplifting"
  ],
  "sleep": [
    "Brooding",
    "Calm",
    "Chill",
    "Dark",
    "Dreamlike",
    "Epic",
    "Floating",
    "Heavy",
    "Meditative",
    "Mysterious",
    "Optimistic",
    "Ponderous",
    "Serene",
    "Strong"
  ],
  "meditate": [
    "Brooding",
    "Calm",
    "Chill",
    "Dark",
    "Downtempo",
    "Dreamlike",
    "Driving",
    "Energizing",
    "Epic",
    "Floating",
    "Heavy",
    "Hopeful",
    "Inspiring",
    "Meditative",
    "Mysterious",
    "Optimistic",
    "Playful",
    "Ponderous",
    "Serene",
    "Strong",
    "Upbeat",
    "Uplifting"
  ]
}
*/

type QueueEntry = {
  url: string;
  folder: string;
  filename: string;
  songId: string;
};

class DownloadQueue {
  private queue: QueueEntry[] = [];
  private activeDownloads = 0;
  private maxConcurrency: number;
  private progressState: ProgressState;

  constructor(maxConcurrency: number, progressState: ProgressState) {
    this.maxConcurrency = maxConcurrency;
    this.progressState = progressState;
  }

  public enqueue(url: string, folder: string, filename: string, songId: string) {
    // Skip if already processed
    if (this.progressState.processedSongs.has(songId)) {
      return;
    }
    this.queue.push({ url, folder, filename, songId });
    this.processQueue();
  }

  public async waitForCompletion(): Promise<void> {
    return new Promise((resolve) => {
      const checkInterval = setInterval(() => {
        if (this.activeDownloads === 0 && this.queue.length === 0) {
          clearInterval(checkInterval);
          resolve();
        }
      }, 100);
    });
  }

  private async processQueue(): Promise<void> {
    while (this.activeDownloads < this.maxConcurrency && this.queue.length > 0) {
      const { url, folder, filename, songId } = this.queue.shift()!;
      this.activeDownloads++;

      this.processDownload(url, folder, filename, songId).then(() => {
        this.activeDownloads--;
        this.processQueue();
      }).catch((error) => {
        console.error('Error downloading song:', error);
        this.activeDownloads--;
        this.processQueue();
      });
    }
  }

  private async processDownload(url: string, folder: string, filename: string, songId: string): Promise<void> {
    await downloadSong(url, folder, filename);
    this.progressState.songCount++;
    this.progressState.processedSongs.add(songId);
    songCount = this.progressState.songCount;
    
    // Save progress every 10 songs
    if (this.progressState.songCount % 10 === 0) {
      saveProgress(this.progressState);
    }
    
    console.log(`${songCount} songs downloaded successfully \n${this.queue.length} remaining`);
  }
}

// Load or initialize progress
const savedProgress = loadProgress();
const progressState: ProgressState = savedProgress || {
  songCount: 0,
  currentMentalStateIndex: 0,
  currentGenreIndex: 0,
  processedSongs: new Set(),
  completedGenres: []
};

songCount = progressState.songCount;

const downloadQueue = new DownloadQueue(3, progressState);

// Set up graceful shutdown
process.on('SIGINT', () => {
  console.log(chalk.yellow('\n\nReceived SIGINT, saving progress...'));
  saveProgress(progressState);
  process.exit(0);
});

process.on('SIGTERM', () => {
  console.log(chalk.yellow('\n\nReceived SIGTERM, saving progress...'));
  saveProgress(progressState);
  process.exit(0);
});

for (let mentalStateIndex = progressState.currentMentalStateIndex; mentalStateIndex < mentalStates.length; mentalStateIndex++) {
  const mentalState = mentalStates[mentalStateIndex];
  console.log(chalk.red(`Starting mental state ${mentalState}`))
  
  const allGenres = [...genres[mentalState].base, ...genres[mentalState].nature];
  const startGenreIndex = mentalStateIndex === progressState.currentMentalStateIndex ? progressState.currentGenreIndex : 0;
  
  for (let genreIndex = startGenreIndex; genreIndex < allGenres.length; genreIndex++) {
    const genre = allGenres[genreIndex];
    const genreKey = `${mentalState}:${genre}`;
    
    // Skip if already completed
    if (progressState.completedGenres.includes(genreKey)) {
      console.log(chalk.gray(`Skipping already completed: ${genre} (${mentalState})`));
      continue;
    }
    
    console.log(chalk.yellow(`Starting genre ${genre}`))
    
    // Update current position
    progressState.currentMentalStateIndex = mentalStateIndex;
    progressState.currentGenreIndex = genreIndex;

    // 
    //  Phase 1 : Fetch all song data
    // 
    const response = await fetch(`https://api.brain.fm/v3/servings/search?genre=${genre}&dynamicMentalStateId=${mentalState}`, {
      headers: {
        authorization: `Bearer ${AUTHTOKEN}`
      }
    });
    if (!response.ok) {
      chalk.red("ERROR SEARCHING SONGS")
      chalk.red(response.text());
      continue;
    }
    const responseData = await response.json();
    const data = formatAudioData(responseData.result);

    if (checkIfJsonExists(`./json-data/${mentalState}/${genre}.json`)) continue;
    ensureDirectory(`./json-data/${mentalState}`);
    let file = fs.createWriteStream(`./json-data/${mentalState}/${genre}.json`);

    file.write(JSON.stringify(data));

    file.close();

    // 
    // Phase 2 : Download songs to device
    // 

    console.log(chalk.green(`Started queuing ${genre} ${mentalState}`))
    for (const song of data) {
      let activity = song.track.tags.filter((x: any) => x.type == 'activity').map((x: any) => x.value).join('/');
      let NEL = song.trackVariation.neuralEffectLevel;
      let level = (NEL > 0.66 ? "high" : NEL > 0.33 ? "medium" : "low");

      let folder = `./songs/${genre}/${mentalState}/${activity}/${level}`
      let filename = song.trackVariation.baseUrl;
      let downloadLink = song.trackVariation.tokenedUrl;
      let songId = `${mentalState}:${genre}:${filename}`;

      downloadQueue.enqueue(downloadLink, folder, filename, songId);
    }
    
    // Wait for all songs in this genre to complete before marking as done
    await downloadQueue.waitForCompletion();
    
    // Mark genre as completed
    progressState.completedGenres.push(genreKey);
    saveProgress(progressState);
    console.log(chalk.green(`✓ Completed ${genre} (${mentalState})`));

  }
}

// All done, clear progress file
clearProgress();
console.log(chalk.green.bold(`\n🎉 All downloads complete! Total songs: ${songCount}`));


function downloadSong(downloadLink: string, folder: string, filename: string): Promise<void> {
  return new Promise((resolve, reject) => {
    ensureDirectory(folder)
    https.get(downloadLink, (response) => {
      response.pipe(fs.createWriteStream(`${folder}/${filename}`))
        .on('finish', () => {
          resolve();
        })
        .on('error', (error: Error) => {
          console.error('Error downloading song:', error);
          reject(error);
        });
    });
  });
}

function formatAudioData(arr: any[]) {
  return arr.map(item => {
    delete item.track.similarTracks;
    return item;
  });
}

function ensureDirectory(directory: string) {
  if (!fs.existsSync(directory)) {
    fs.mkdirSync(directory, { recursive: true });
  }
}

function checkIfJsonExists(filePath: string) {
  return fs.existsSync(filePath);
}