#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;
	}
	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++) {
				memset(this->tmpBuffer, 0x00, DVDCSS_BLOCK_SIZE);
				int dataRd = dvdcss_read(driveIoCtl->GetDvdCssHandle(), this->tmpBuffer, 1, (this->decrypt ? DVDCSS_READ_DECRYPT : DVDCSS_NOFLAGS));
				if (dataRd < 0) {
					dvdcss_seek(driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar, (this->decrypt ? DVDCSS_SEEK_KEY : DVDCSS_NOFLAGS) );

					if ((retry + 1) >= numRetry) {
						dvdcss_seek(driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar + 1, (this->decrypt ? DVDCSS_SEEK_KEY : DVDCSS_NOFLAGS));
						this->failedReads++;
					}

					this->retriedReads++;
					continue; // retry
				}
				else {
					this->iso->write((const char*)this->tmpBuffer, DVDCSS_BLOCK_SIZE * dataRd);
					break;
				}
			}
			nmRd += 1;
			this->sectorsReadSoFar += 1;
			this->fileReadSoFar += 1;
			
			if (this->sectorsReadSoFar > this->drive->Sectors()) // check if we've read the last sector ... 
				break;
		}

		return nmRd;

	}

	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 (dvdcss_was_error(driveIoCtl->GetDvdCssHandle())) {
				this->errorMsg = std::string(dvdcss_error(driveIoCtl->GetDvdCssHandle()));
			}
			if (numRead < 0) {
				dvdcss_seek(driveIoCtl->GetDvdCssHandle(), this->sectorsReadSoFar, DVDCSS_SEEK_KEY);
				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->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->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->sectorsReadSoFar = 0;
		this->fileReadSoFar = 0;
		this->keepRunning = false;
		this->ripinThread->join();

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

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

		this->iso = nullptr;
	}
}