parser.cpp

#include "parser.h"
#include "BinaryReader.h"
#include <iostream>
#include <fstream>

OsuDbParser::OsuDbParser() : dbInfo{0} {}

OsuDbParser::~OsuDbParser() {}

bool OsuDbParser::Load(const std::string& path) {
    std::ifstream fs(path, std::ios::binary | std::ios::ate);
    if (!fs.is_open()) {
        std::cerr << "Failed to open " << path << std::endl;
        return false;
    }

    std::streamsize size = fs.tellg();
    fs.seekg(0, std::ios::beg);

    std::vector<unsigned char> buffer(size);
    if (!fs.read(reinterpret_cast<char*>(buffer.data()), size)) {
        std::cerr << "Failed to read file content" << std::endl;
        return false;
    }

    BinaryReader reader(buffer);
    std::cout << "Buffer size: " << buffer.size() << std::endl;

    try {
        dbInfo.version = reader.ReadInt();
        dbInfo.folderCount = reader.ReadInt();
        dbInfo.accountUnlocked = reader.ReadBool();
        dbInfo.accountUnlockDate = reader.ReadLong();
        dbInfo.playerName = reader.ReadString();
        dbInfo.beatmapCount = reader.ReadInt();
        
        std::cout << "DB Version: " << dbInfo.version << std::endl;
        std::cout << "Beatmaps to parse: " << dbInfo.beatmapCount << std::endl;
        dbInfo.beatmaps.reserve(dbInfo.beatmapCount);

        for (int i = 0; i < dbInfo.beatmapCount; ++i) {
            try {
                BeatmapEntry entry = ParseBeatmap(reader);
                
                dbInfo.beatmaps.push_back(entry);
                existingBeatmapIds.insert(entry.beatmapId);
                existingSetIds.insert(entry.threadId);

                if (i % 1000 == 0) std::cout << "Parsed " << i << "/" << dbInfo.beatmapCount << " (BID: " << entry.beatmapId << ")" << std::endl;
            
            } catch (const std::exception& e) {
                std::cerr << "Error parsing beatmap " << i << ": " << e.what() << " at pos " << reader.Tell() << std::endl;
                std::cerr << "Attempting to recover..." << std::endl;
                FindNextBeatmap(reader);
            }
        }
        
        std::cout << "Finished parsing. Total Entries: " << dbInfo.beatmaps.size() << std::endl;

    } catch (const std::exception& e) {
        std::cerr << "Error parsing DB: " << e.what() << " at pos " << reader.Tell() << std::endl;
        return false;
    }
    return true;
}

BeatmapEntry OsuDbParser::ParseBeatmap(BinaryReader& reader) {
    BeatmapEntry entry;
    
    if (reader.Tell() >= reader.GetLength()) {
        throw std::runtime_error("Unexpected end of stream");
    }

    entry.artistName = reader.ReadString();
    entry.artistNameUnicode = reader.ReadString();
    entry.songTitle = reader.ReadString();
    entry.songTitleUnicode = reader.ReadString();
    entry.creatorName = reader.ReadString();
    entry.difficulty = reader.ReadString();
    entry.audioFileName = reader.ReadString();
    entry.md5Hash = reader.ReadString();
    entry.osuFileName = reader.ReadString();
    
    entry.rankedStatus = reader.ReadByte();
    entry.hitCircleCount = reader.ReadShort();
    entry.sliderCount = reader.ReadShort();
    entry.spinnerCount = reader.ReadShort();
    entry.lastModificationTime = reader.ReadLong();
    entry.approachRate = reader.ReadFloat();
    entry.circleSize = reader.ReadFloat();
    entry.hpDrain = reader.ReadFloat();
    entry.overallDifficulty = reader.ReadFloat();
    entry.sliderVelocity = reader.ReadDouble();
    
    // Star Ratings
    auto readStarRating = [&](std::vector<std::pair<int, double>>& vec, const char* name) {
        int count = reader.ReadInt();
        for (int j = 0; j < count; ++j) {
            unsigned char b1 = reader.ReadByte(); // 0x08
            if (b1 != 0x08) {
                 // std::cerr << "Warning: Expected 0x08, got " << std::hex << (int)b1 << std::dec << " at " << reader.Tell()-1 << std::endl;
            }
            
            int mod = reader.ReadInt();
            unsigned char type = reader.ReadByte();
            double rating = 0;
            
            if (type == 0x0B) rating = reader.ReadDouble(); // Some docs say 0x0B
            else if (type == 0x0C) rating = reader.ReadFloat(); // Float
            else if (type == 0x0D) rating = reader.ReadDouble(); // Double (Standard)
            else {
                 std::string err = "Unknown star rating value type: " + std::to_string((int)type) + " at pos " + std::to_string(reader.Tell()-1);
                 throw std::runtime_error(err);
            }
            vec.push_back({mod, rating});
        }
    };

    readStarRating(entry.starRatingStd, "Std");
    readStarRating(entry.starRatingTaiko, "Taiko");
    readStarRating(entry.starRatingCTB, "CTB");
    readStarRating(entry.starRatingMania, "Mania");
    
    entry.drainTime = reader.ReadInt();
    entry.totalTime = reader.ReadInt();
    entry.audioPreviewTime = reader.ReadInt();
    
    // Timing Points
    int timingPointCount = reader.ReadInt();
    for(int j=0; j<timingPointCount; ++j) {
        TimingPoint tp;
        tp.bpm = reader.ReadDouble();
        tp.offset = reader.ReadDouble();
        tp.inherited = reader.ReadBool();
        entry.timingPoints.push_back(tp);
    }
    
    entry.difficultyId = reader.ReadInt();
    entry.beatmapId = reader.ReadInt();
    entry.threadId = reader.ReadInt();
    
    entry.standardGrade = reader.ReadByte();
    entry.taikoGrade = reader.ReadByte();
    entry.ctbGrade = reader.ReadByte();
    entry.maniaGrade = reader.ReadByte();
    
    entry.localOffset = reader.ReadShort();
    entry.stackLeniency = reader.ReadFloat();
    entry.gameplayMode = reader.ReadByte();
    
    entry.songSource = reader.ReadString();
    entry.songTags = reader.ReadString();
    entry.onlineOffset = reader.ReadShort();
    entry.font = reader.ReadString();
    entry.isUnplayed = reader.ReadBool();
    entry.lastPlayed = reader.ReadLong();
    entry.isOsz2 = reader.ReadBool();
    entry.folderName = reader.ReadString();
    entry.lastCheckedAgainstOsuRepo = reader.ReadLong();
    entry.ignoreBeatmapSound = reader.ReadBool();
    entry.ignoreBeatmapSkin = reader.ReadBool();
    entry.disableStoryboard = reader.ReadBool();
    entry.disableVideo = reader.ReadBool();
    entry.visualOverride = reader.ReadBool();
    
    // Handle optional/variable fields
    if (dbInfo.version >= 20140609) {
        entry.lastModificationTime2 = reader.ReadInt();
        entry.maniaScrollSpeed = reader.ReadByte();
    }

    return entry;
}

void OsuDbParser::FindNextBeatmap(BinaryReader& reader) {
    // Heuristic: Scan for a sequence that looks like the start of a beatmap.
    // Start of beatmap:
    // String (Artist)
    // String (ArtistUnicode)
    // String (Title)
    // String (TitleUnicode)
    // String (Creator)
    // String (Difficulty)
    // String (AudioFile)
    // String (MD5)
    // String (OsuFile)
    // Byte (RankedStatus)
    
    // We look for 9 valid strings followed by a byte that is a valid RankedStatus (0-7 usually).
    
    size_t startPos = reader.Tell();
    size_t bufferSize = reader.GetLength();
    const unsigned char* data = reader.GetData();
    
    // Limit scan to avoid hanging
    size_t scanLimit = 10000000; // 10MB scan limit
    
    for (size_t i = 0; i < scanLimit && (startPos + i) < bufferSize; ++i) {
        size_t currentPos = startPos + i;
        reader.Seek(currentPos);
        
        bool match = true;
        try {
            // Try to read 9 strings
            for (int k = 0; k < 9; ++k) {
                unsigned char b = reader.PeekByte();
                if (b == 0x00) {
                    reader.ReadByte(); // Empty string
                } else if (b == 0x0B) {
                    reader.ReadString(); // Normal string
                } else {
                    match = false;
                    break;
                }
            }
            
            if (match) {
                // Check RankedStatus
                if (reader.Tell() < bufferSize) {
                    unsigned char status = reader.ReadByte();
                    if (status <= 7) { // Valid ranked status range roughly
                         // Found it!
                         std::cout << "Recovered at offset " << currentPos << std::endl;
                         reader.Seek(currentPos);
                         return;
                    }
                }
            }
        } catch (...) {
            match = false;
        }
        
        // If we threw or didn't match, continue scanning
    }
    
    std::cerr << "Failed to find next beatmap. Aborting." << std::endl;
    // Seek to end to stop loop
    reader.Seek(bufferSize);
}