shimmeringmoon/src/arcaea/import_charts.rs

use std::collections::HashMap;

use anyhow::{anyhow, Context};
use serde::Deserialize;

use crate::{
	arcaea::{chart::Side, rating::rating_as_fixed},
	context::paths::ShimmeringPaths,
};

use super::{
	chart::{Difficulty, Level},
	rating::{rating_from_fixed, Rating},
};

// {{{ Notecount
struct NotecountEntry {
	difficulty: Difficulty,
	level: Level,
	name: String,
	notecount: u32,
}

fn get_notecount_records(paths: &ShimmeringPaths) -> anyhow::Result<Vec<NotecountEntry>> {
	let mut entries = Vec::new();
	let mut reader = csv::Reader::from_reader(std::io::BufReader::new(std::fs::File::open(
		paths.notecount_path(),
	)?));

	for result in reader.records() {
		let record = result?;

		let notecount = record
			.get(0)
			.ok_or_else(|| anyhow!("Missing notecount in csv entry"))?
			.parse()?;

		let raw_difficulty = record
			.get(1)
			.ok_or_else(|| anyhow!("Missing level/difficulty in csv entry"))?;

		let name = record
			.get(2)
			.ok_or_else(|| anyhow!("Missing name in csv entry"))?;

		let (raw_difficulty, raw_level) = raw_difficulty
			.split_once(" ")
			.ok_or_else(|| anyhow!("Invalid level/difficulty string in csv entry"))?;

		entries.push(NotecountEntry {
			notecount,
			name: name.to_owned(),
			level: raw_level.parse()?,
			difficulty: raw_difficulty.parse()?,
		});
	}

	Ok(entries)
}
// }}}
// {{{ PTT entries
#[derive(Clone, Copy, Deserialize)]
struct PTTEntry {
	#[serde(rename = "0")]
	pst: Option<f32>,
	#[serde(rename = "1")]
	prs: Option<f32>,
	#[serde(rename = "2")]
	ftr: Option<f32>,
	#[serde(rename = "3")]
	byd: Option<f32>,
	#[serde(rename = "4")]
	etr: Option<f32>,
}

impl PTTEntry {
	fn get_rating(&self, difficulty: Difficulty) -> Option<Rating> {
		let float = match difficulty {
			Difficulty::PST => self.pst,
			Difficulty::PRS => self.prs,
			Difficulty::FTR => self.ftr,
			Difficulty::BYD => self.byd,
			Difficulty::ETR => self.etr,
		};

		float.map(|f| rating_from_fixed((f * 100.0).round() as i32))
	}
}

fn get_ptt_entries(paths: &ShimmeringPaths) -> anyhow::Result<HashMap<String, PTTEntry>> {
	let result = serde_json::from_reader(std::io::BufReader::new(std::fs::File::open(
		paths.cc_data_path(),
	)?))?;

	Ok(result)
}
// }}}
// {{{ Songlist types
#[derive(Deserialize)]
struct LocalizedName {
	en: String,
	og: Option<String>,
}

impl LocalizedName {
	fn get(&self) -> &str {
		self.og.as_ref().unwrap_or(&self.en)
	}
}

#[derive(Deserialize)]
struct Chart {
	rating: u8,
	#[serde(default, rename = "ratingPlus")]
	rating_plus: bool,
	#[serde(rename = "ratingClass")]
	difficulty: u8,

	#[serde(rename = "chartDesigner")]
	chart_designer: String,

	#[allow(unused)]
	#[serde(rename = "jacketDesigner")]
	jacket_designer: String,

	#[serde(rename = "title_localized")]
	title: Option<LocalizedName>,
}

#[derive(Deserialize)]
struct Song {
	#[serde(rename = "idx")]
	id: u32,
	#[serde(rename = "id")]
	shorthand: String,
	#[serde(rename = "title_localized")]
	title: LocalizedName,

	artist: String,
	bpm: String,
	side: u32,
	difficulties: Vec<Chart>,
}

#[derive(Deserialize)]
struct DeletedSong {
	#[allow(unused)]
	deleted: bool,
}

#[derive(Deserialize)]
#[serde(untagged)]
enum SonglistEntry {
	Song(Song),

	#[allow(unused)]
	Deleted(DeletedSong),
}

#[derive(Deserialize)]
struct Songlist {
	songs: Vec<SonglistEntry>,
}
// }}}
// {{{ Process songlist file
pub fn import_songlist(
	paths: &ShimmeringPaths,
	conn: &mut rusqlite::Connection,
) -> anyhow::Result<()> {
	let notecount_records =
		get_notecount_records(paths).context("Failed to read notecount records")?;
	let ptt_entries = get_ptt_entries(paths).context("Failed to read ptt entries")?;

	let transaction = conn.transaction()?;
	transaction.execute("DELETE FROM charts", ())?;
	transaction.execute("DELETE FROM songs", ())?;

	let songlist: Songlist = serde_json::from_reader(std::io::BufReader::new(
		std::fs::File::open(paths.songlist_path())?,
	))?;

	let mut song_count = 0;
	let mut chart_count = 0;

	for song in songlist.songs {
		let song = match song {
			SonglistEntry::Song(song) => song,
			SonglistEntry::Deleted(_) => continue,
		};

		song_count += 1;
		transaction.execute(
			"
        INSERT INTO songs(id,title,shorthand,artist,side,bpm)
        VALUES (?,?,?,?,?,?)
      ",
			(
				song.id,
				song.title.get(),
				&song.shorthand,
				&song.artist,
				Side::SIDES[song.side as usize],
				song.bpm,
			),
		)?;

		for chart in song.difficulties {
			if chart.rating == 0 {
				continue;
			}

			chart_count += 1;

			let difficulty = crate::private_server::decode_difficulty(chart.difficulty)
				.ok_or_else(|| anyhow!("Invalid difficulty"))?;

			let level = format!(
				"{}{}",
				chart.rating,
				if chart.rating_plus { "+" } else { "" }
			)
			.parse()
			.context("Failed to parse level")?;

			let name = chart.title.as_ref().unwrap_or(&song.title).get();
			let notecount = notecount_records
				.iter()
				.find_map(|record| {
					let names_match = record.name == name
						|| record.name == format!("{name} ({})", &song.artist)
						|| record.name == song.shorthand;

					if names_match && record.level == level && record.difficulty == difficulty {
						Some(record.notecount)
					} else {
						None
					}
				})
				.ok_or_else(|| {
					anyhow!(
						"Cannot find note count for song '{}' [{}]",
						name,
						difficulty
					)
				})?;

			let cc = ptt_entries
				.get(&song.shorthand)
				.ok_or_else(|| anyhow!("Cannot find PTT data for song '{}'", song.shorthand))?
				.get_rating(difficulty)
				.ok_or_else(|| {
					anyhow!("Cannot find PTT data for song '{}' [{}]", name, difficulty)
				})?;

			transaction.execute(
				"
          INSERT INTO charts(
            song_id, title, difficulty,
            level, note_count, chart_constant,
            note_design
          ) VALUES(?,?,?,?,?,?,?)
        ",
				(
					song.id,
					chart.title.as_ref().map(|t| t.get()),
					difficulty,
					level,
					notecount,
					rating_as_fixed(cc),
					chart.chart_designer,
				),
			)?;
		}
	}

	transaction.commit()?;

	println!("✅ Succesfully imported {chart_count} charts, {song_count} songs");

	Ok(())
}
// }}}