#include "DvdRipper.hpp"
#include <dvdcss.h>

#include <iostream>
#include <fstream>
#include <thread>
#include <sstream>
#include <filesystem>

#include "../Scsi/IoCtl.hpp"
#include "../Scsi/OpticalDrive.hpp"
#include "../Dvd/TitleKey.hpp"

namespace Li::Dvd {

	uint32_t DvdRipper::FailedReads() {
		return this->failedReads;
	}
	uint32_t DvdRipper::RetriedReads() {
		return this->retriedReads;
	}
	uint32_t DvdRipper::FileSoFar() {
		return this->fileReadSoFar;
	}
	uint32_t DvdRipper::FileLen() {
		return this->fileRemain;
	}
	float DvdRipper::PercentDone() {
		return (((float)this->sectorsReadSoFar / (float)this->drive->Sectors()));
	}
	uint32_t DvdRipper::SectorsReadSoFar() {
		return this->sectorsReadSoFar;
	}
	std::string DvdRipper::OutputPath() {
		return this->outputPath;
	}
	std::byte* DvdRipper::GetDiscKey() {
		return this->drive->DiscKey();
	}
	std::byte* DvdRipper::GetTitleKey() {
		return this->titleKey;
	}
	bool DvdRipper::Decrypt() {
		return this->decrypt;
	}
	bool DvdRipper::Starting() {
		return this->starting;
	}
	bool DvdRipper::Done() {
		return this->done;
	}
	bool DvdRipper::Error() {
		return this->error;
	}
	std::string DvdRipper::ErrorMessage() {
		return this->errorMsg;
	}

	// Function to handle bad sectors
	// the way this algorithm essentially works is as follows:
	// Attempt to read total num of sectors, one at a time
	// 
	// - if total read sectors count is < 0 (i.e read failed), 
	// - check if the last read was also an error, if it was not, 
	// - try to read the sector again, retry it 5 times
	// -- if it still fails, set that sector to all 0x00, and skip it
	// --- also set a flag saying the last sector read failed, keep skipping until get a successful read,
	int DvdRipper::read1SectorATimeSkippingBadSectors(int toRead) {		
		int nmRd = 0;
		int numRetry = 5;
		for (int i = 0; i < toRead; i++) {
			for (int retry = 0; retry < numRetry; retry++) 
			{
				int dataRd = dvdcss_read(driveIoCtl->GetDvdCssHandle(), this->tmpBuffer, 1, (this->decrypt ? DVDCSS_READ_DECRYPT : DVDCSS_NOFLAGS));
				if (dataRd < 0 || dvdcss_was_error(driveIoCtl->GetDvdCssHandle())) {
					
					// if the last read was an error, and this read is an error, give up trying as were still in 
					// the section of bad sectors,
					
					// otherwise retry a few times, to see if we can read this sector
					if (this->lastReadWasErr || ((retry + 1) >= numRetry)) { 
						this->errorMsg = "one or more sectors failed to be read.";

						// try seek forward 1 sector
						dvdcss_seek(driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar + 1, DVDCSS_NOFLAGS);
						
						// set the last read was error flag,
						// increment the number of failed reads
						// and set the error flag to trigger the "This ISO Was created with errors" message.
						this->failedReads++;
						this->lastReadWasErr = true;
						this->error = true;

						// set current sector to all 0x00 ...
						memset(this->tmpBuffer, 0x00, DVDCSS_BLOCK_SIZE);

						// write null sector to the ISO ...
						this->iso->write((const char*)this->tmpBuffer, DVDCSS_BLOCK_SIZE);
						break;
					}
					else {
						// if we have not reached max retry count, then set the position back to the current sector
						// and attempt a read again ..
						dvdcss_seek(driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar, DVDCSS_NOFLAGS);
					}

					this->retriedReads++;
				}
				else {
					// if there was a successful read, 
					this->lastReadWasErr = false;
					this->iso->write((const char*)this->tmpBuffer, DVDCSS_BLOCK_SIZE * dataRd);
					break;
				}
			}
			// increment the num sectors read counter
			nmRd += 1;
			this->sectorsReadSoFar += 1;
			this->fileReadSoFar += 1;

			// check if we've read the last sector ...
			if (this->sectorsReadSoFar > this->drive->Sectors())  
				break;
		}

		return nmRd;
	}

	// this function is main function for reading & decrypting sectors
	// on a Content Scramble System enabled DVD.
	void DvdRipper::readCssSectors() {
		do {
			int toRead = this->sectorsAtOnce;
			if ((this->fileReadSoFar + toRead) > this->fileRemain) {
				toRead = (this->fileRemain - this->fileReadSoFar);
			}

			int numRead = dvdcss_read(driveIoCtl->GetDvdCssHandle(), this->tmpBuffer, toRead, DVDCSS_READ_DECRYPT);

			if (numRead < 0 || dvdcss_was_error(driveIoCtl->GetDvdCssHandle())) {
				dvdcss_seek(driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar, DVDCSS_NOFLAGS);
				read1SectorATimeSkippingBadSectors(toRead);
			}
			else {
				this->iso->write((const char*)this->tmpBuffer, DVDCSS_BLOCK_SIZE * numRead);
				this->sectorsReadSoFar += numRead;
				this->fileReadSoFar += numRead;
			}
		} while ( (this->fileReadSoFar < this->fileRemain) && 
				  (this->sectorsReadSoFar < this->drive->Sectors()) && 
			      this->keepRunning );

		this->inFile = Li::Dvd::TitleKey::IsSectorInFile(this->sectorsReadSoFar);
		this->fileReadSoFar = 0;
		this->fileRemain = 0;
	}
	void DvdRipper::readNonFileSectors() {
		dvdcss_title(this->driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar);
		this->fileRemain = Li::Dvd::TitleKey::GetDistanceToNextFile(this->sectorsReadSoFar);

		if (this->fileRemain == 0 || (this->sectorsReadSoFar + this->fileRemain) > this->drive->Sectors()) {
			this->fileRemain = (this->drive->Sectors() - this->sectorsReadSoFar);
		}

		this->fileReadSoFar = 0;
		this->readCssSectors();
	}

	void DvdRipper::readFileSectors() {
		dvdcss_title(this->driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar);
		this->titleKey = (std::byte*)dvdcss_get_cur_titlekey(this->driveIoCtl->GetDvdCssHandle());

		sectorInfo* sectorInf = Li::Dvd::TitleKey::GetSectorInfo(this->sectorsReadSoFar);
		this->fileRemain = (sectorInf->endSector - sectorInf->startSector);

		if ((this->sectorsReadSoFar + this->fileRemain) > this->drive->Sectors()) {
			this->fileRemain = (this->drive->Sectors() - this->sectorsReadSoFar);
		}

		this->fileReadSoFar = 0;

		this->readCssSectors();
	}

	void DvdRipper::readUnprotectedSectors() {
		do {
			int toRead = this->sectorsAtOnce;
			if ((this->sectorsReadSoFar + toRead) > this->drive->Sectors()) {
				toRead = this->drive->Sectors() - this->sectorsReadSoFar;
			}

			int numRead = dvdcss_read(this->driveIoCtl->GetDvdCssHandle(), this->tmpBuffer, toRead, DVDCSS_NOFLAGS);

			if (dvdcss_was_error(driveIoCtl->GetDvdCssHandle())) {
				this->errorMsg = std::string(dvdcss_error(driveIoCtl->GetDvdCssHandle()));
			}
			if (numRead < 0) { // Handle bad sectors
				dvdcss_seek(driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar, DVDCSS_NOFLAGS);
				read1SectorATimeSkippingBadSectors(toRead);
			}
			else {
				this->iso->write((const char*)this->tmpBuffer, DVDCSS_BLOCK_SIZE * numRead);
				this->sectorsReadSoFar += numRead;
				this->fileReadSoFar += numRead;
			}

		} while (this->sectorsReadSoFar < this->drive->Sectors() && this->keepRunning);
	}

	void DvdRipper::ripThread() {
		this->sectorsReadSoFar = 0;

		// due to how libdvdcss works, you can only actually decrypt a dvd sector
		// if you seek directly at it or call dvdcss_title on the sector containing the start of the encrypted title file
		// so, i have to actually parse the disc file system, in order to find all those
		if (this->decrypt) {
			Li::Dvd::TitleKey::FindTitleKeys(driveIoCtl->GetDvdCssHandle());
			dvdcss_seek(driveIoCtl->GetDvdCssHandle(), 0, DVDCSS_SEEK_KEY);
			this->inFile = Li::Dvd::TitleKey::IsSectorInFile(this->sectorsReadSoFar);

			dvdcss_title(this->driveIoCtl->GetDvdCssHandle(), 0);
			this->titleKey = (std::byte*)dvdcss_get_cur_titlekey(this->driveIoCtl->GetDvdCssHandle());
		}

		this->starting = false;

		do {
			if (this->decrypt) {
				if (this->inFile) {
					// read & decrypt all files (including those associated w the filesystem)
					this->readFileSectors();
				}
				else {
					// read the data not associated with any file on the filesystem.
					this->readNonFileSectors();
				}
			}
			else {
				// for non scrambled dvds, just read every sector one after another.
				this->readUnprotectedSectors();
			}


		} while (this->sectorsReadSoFar < this->drive->Sectors() && this->keepRunning);

		this->done = true;
	}

	DvdRipper::DvdRipper(Li::Scsi::OpticalDrive* opticalDrive, std::string outputISO, uint32_t sectorsAtOnce) {
		this->drive = opticalDrive;
		this->outputPath = outputISO;
		this->sectorsAtOnce = sectorsAtOnce;

		this->failedReads = 0;
		this->retriedReads = 0;

		this->sectorsReadSoFar = 0;
		this->fileReadSoFar = 0;
		this->fileRemain = 0;

		this->inFile = false;
		this->done = false;
		this->error = false;
		this->errorMsg = "";
		this->keepRunning = false;
		this->starting = false;
		this->lastReadWasErr = false;

		this->ripinThread = nullptr;
		this->tmpBuffer = new std::byte[DVDCSS_BLOCK_SIZE * this->sectorsAtOnce];

	}

	DvdRipper::~DvdRipper() {
		delete this->tmpBuffer;

		if (this->ripinThread != nullptr) {
			this->EndRip();
		}
		if (this->iso != nullptr) {
			if (this->iso->is_open())
				this->iso->close();
			delete this->iso;
		}
	}

	void DvdRipper::StartRip(uint32_t driveSpeed, bool cssDecrypt) {

		this->driveIoCtl = new Li::Scsi::IoCtl(this->drive->DrivePath());
		this->driveIoCtl->AllowReadingPastDisc();
		this->driveIoCtl->SetDriveSpeed(driveSpeed, driveSpeed);

		this->decrypt = cssDecrypt;
		if (!dvdcss_is_scrambled(this->driveIoCtl->GetDvdCssHandle())) {
			this->decrypt = false;
		}
		
		this->starting = true;
		this->iso = new std::ofstream(this->outputPath, std::ios::binary);
		this->keepRunning = true;

		//this->driveIoCtl->ExclusiveLockDrive();
		this->ripinThread = new std::thread(&DvdRipper::ripThread, this);

	}

	void DvdRipper::EndRip() {
		this->done = false;
		this->error = false;
		this->errorMsg = "";

		this->inFile = false;
		this->keepRunning = false;
		this->starting = false;
		this->lastReadWasErr = false;

		this->sectorsReadSoFar = 0;
		this->fileReadSoFar = 0;
		this->ripinThread->join();

		delete this->driveIoCtl;
		delete this->ripinThread;

		this->ripinThread = nullptr;
		iso->close();
		delete iso;

		this->iso = nullptr;
	}
}