#include "baseaddress.h"
#include <QScrollArea>
#include <QScrollBar>

using namespace std;


BNBaseAddressDetectionPOISetting BaseAddressDetectionPOISettingFromString(const string& setting)
{
	if (setting == "Strings only")
		return POIAnalysisStringsOnly;
	if (setting == "Functions only")
		return POIAnalysisFunctionsOnly;
	return POIAnalysisAll; // Default to All
}


string BaseAddressDetectionConfidenceToString(BNBaseAddressDetectionConfidence level)
{
	switch (level)
	{
		case NoConfidence:
			return "Unassigned";
		case HighConfidence:
			return "High";
		case LowConfidence:
			return "Low";
		default:
			return "Unknown";
	}
}


BaseAddressDetectionThread::BaseAddressDetectionThread(BaseAddressDetectionQtInputs* widgetInputs,
	BinaryNinja::Ref<BinaryNinja::BinaryView> bv)
{
	m_inputs = widgetInputs;
	m_view = bv;
	m_baseDetection = new BinaryNinja::BaseAddressDetection(m_view);
}


BaseAddressDetectionThread::~BaseAddressDetectionThread()
{
	if (m_baseDetection)
		delete m_baseDetection;
}


void BaseAddressDetectionThread::run()
{
	BaseAddressDetectionQtResults results;
	uint64_t value;
	string errorStr;

	if (!BinaryNinja::BinaryView::ParseExpression(
		m_view, m_inputs->AlignmentLineEdit->text().toStdString(), value, 0, errorStr))
	{
		results.Status = "Invalid alignment value (" + errorStr + ")";
		emit ResultReady(results);
		return;
	}
	uint32_t alignment = value;

	if (!BinaryNinja::BinaryView::ParseExpression(
		m_view, m_inputs->StrlenLineEdit->text().toStdString(), value, 0, errorStr))
	{
		results.Status = "Invalid minimum string length (" + errorStr + ")";
		emit ResultReady(results);
		return;
	}
	uint32_t minStrlen = value;

	uint64_t upperBoundary;
	if (!BinaryNinja::BinaryView::ParseExpression(
		m_view, m_inputs->UpperBoundary->text().toStdString(), upperBoundary, 0, errorStr))
	{
		results.Status = "Invalid upper boundary address (" + errorStr + ")";
		emit ResultReady(results);
		return;
	}

	uint64_t lowerBoundary;
	if (!BinaryNinja::BinaryView::ParseExpression(
		m_view, m_inputs->LowerBoundary->text().toStdString(), lowerBoundary, 0, errorStr))
	{
		results.Status = "Invalid lower boundary address (" + errorStr + ")";
		emit ResultReady(results);
		return;
	}

	if (lowerBoundary >= upperBoundary)
	{
		results.Status = "Upper boundary address is less than lower";
		emit ResultReady(results);
		return;
	}

	if (!BinaryNinja::BinaryView::ParseExpression(
		m_view, m_inputs->MaxPointersPerCluster->text().toStdString(), value, 0, errorStr))
	{
		results.Status = "Invalid max pointers (" + errorStr + ")";
		emit ResultReady(results);
		return;
	}

	uint32_t maxPointersPerCluster = value;
	if (maxPointersPerCluster < 2)
	{
		results.Status = "Invalid max pointers (must be >= 2)";
		emit ResultReady(results);
		return;
	}

	BNBaseAddressDetectionPOISetting poiSetting = BaseAddressDetectionPOISettingFromString(
		m_inputs->POIBox->currentText().toStdString());
	BinaryNinja::BaseAddressDetectionSettings settings = {
		m_inputs->ArchitectureBox->currentText().toStdString(),
		m_inputs->AnalysisBox->currentText().toStdString(),
		minStrlen,
		alignment,
		lowerBoundary,
		upperBoundary,
		poiSetting,
		maxPointersPerCluster,
	};

	if (!m_baseDetection->DetectBaseAddress(settings))
		emit ResultReady(results);

	auto scores = m_baseDetection->GetScores(&results.Confidence, &results.LastTestedBaseAddress);
	results.Scores = scores;
	for (const auto& score : scores)
	{
		auto reasons = m_baseDetection->GetReasonsForBaseAddress(score.second);
		results.Reasons[score.second] = reasons;
	}

	emit ResultReady(results);
}


void BaseAddressDetectionWidget::HideResultsWidgets(bool hide)
{
	if (hide)
	{
		m_preferredBaseLabel->setHidden(true);
		m_preferredBase->setHidden(true);
		m_confidenceLabel->setHidden(true);
		m_confidence->setHidden(true);
		m_resultsTableWidget->setHidden(true);
		m_reloadBase->setHidden(true);
		m_rebaseButton->setHidden(true);
	}
	else
	{
		m_preferredBaseLabel->setHidden(false);
		m_preferredBase->setHidden(false);
		m_confidenceLabel->setHidden(false);
		m_confidence->setHidden(false);
		m_resultsTableWidget->setHidden(false);
		m_reloadBase->setHidden(false);
		m_rebaseButton->setHidden(false);
	}
}


void BaseAddressDetectionWidget::GetClickedBaseAddress(const QModelIndex& index)
{
	if (index.isValid())
	{
		auto baseAddress = m_resultsTableWidget->item(index.row(), 0)->text();
		m_reloadBase->setText(baseAddress);
	}
}


void BaseAddressDetectionWidget::HandleResults(const BaseAddressDetectionQtResults& results)
{
	// Save scroll position to prevent unwanted scrolling when hiding/showing widgets
	QScrollArea* scrollArea = nullptr;
	int savedVerticalScrollPosition = 0;
	int savedHorizontalScrollPosition = 0;
	QWidget* ancestor = parentWidget();
	while (ancestor)
	{
		scrollArea = qobject_cast<QScrollArea*>(ancestor);
		if (scrollArea)
		{
			savedVerticalScrollPosition = scrollArea->verticalScrollBar()->value();
			savedHorizontalScrollPosition = scrollArea->horizontalScrollBar()->value();
			break;
		}
		ancestor = ancestor->parentWidget();
	}

	if (!results.Status.empty())
		m_status->setText(QString::fromStdString(results.Status));

	if (results.Status.empty() && m_worker->IsAborted())
		m_status->setText(QString("Aborted by user (Last Base: 0x%1)").arg(results.LastTestedBaseAddress, 0, 16));

	if (results.Scores.empty())
	{
		if (!m_worker->IsAborted() && results.Status.empty())
			m_status->setText("Completed with no results");
		m_preferredBase->setText("Not available");
		m_confidence->setText("Not available");
	}
	else
	{
		HideResultsWidgets(false);
		if (results.Status.empty() && !m_worker->IsAborted())
			m_status->setText("Completed with results");
		m_preferredBase->setText(QString("0x%1").arg(results.Scores.rbegin()->second, 0, 16));
		m_confidence->setText(QString("%1 (Score: %2)").arg(
			QString::fromStdString(BaseAddressDetectionConfidenceToString(results.Confidence)),
			QString::number(results.Scores.rbegin()->first)));
		m_reloadBase->setText(QString("0x%1").arg(results.Scores.rbegin()->second, 0, 16));
	}

	m_resultsTableWidget->clearContents();
	m_resultsTableWidget->setRowCount(results.Scores.size());
	size_t row = 0;
	for (auto rit = results.Scores.rbegin(); rit != results.Scores.rend(); rit++)
	{
		auto [score, baseaddr] = *rit;
		size_t strHits = 0;
		size_t funcHits = 0;
		size_t dataHits = 0;
		for (const auto& reason : results.Reasons.at(baseaddr))
		{
			switch (reason.POIType)
			{
				case POIString:
					strHits++;
					break;
				case POIFunction:
					funcHits++;
					break;
				case POIDataVariable:
					dataHits++;
					break;
				default:
					break;
			}
		}

		m_resultsTableWidget->setItem(row, 0, new QTableWidgetItem(QString("0x%1").arg(baseaddr, 0, 16)));
		m_resultsTableWidget->setItem(row, 1, new QTableWidgetItem(QString::number(score)));
		m_resultsTableWidget->setItem(row, 2, new QTableWidgetItem(QString::number(strHits)));
		m_resultsTableWidget->setItem(row, 3, new QTableWidgetItem(QString::number(funcHits)));
		m_resultsTableWidget->setItem(row, 4, new QTableWidgetItem(QString::number(dataHits)));
		row++;
	}

	m_abortButton->setHidden(true);
	m_startButton->setHidden(false);
	m_startButton->setEnabled(true);

	// Restore scroll position after widget visibility changes
	if (scrollArea)
	{
		scrollArea->verticalScrollBar()->setValue(savedVerticalScrollPosition);
		scrollArea->horizontalScrollBar()->setValue(savedHorizontalScrollPosition);
	}
}


void BaseAddressDetectionWidget::DetectBaseAddress()
{
	// Save scroll position to prevent unwanted scrolling when hiding/showing widgets
	QScrollArea* scrollArea = nullptr;
	int savedVerticalScrollPosition = 0;
	int savedHorizontalScrollPosition = 0;
	QWidget* ancestor = parentWidget();
	while (ancestor)
	{
		scrollArea = qobject_cast<QScrollArea*>(ancestor);
		if (scrollArea)
		{
			savedVerticalScrollPosition = scrollArea->verticalScrollBar()->value();
			savedHorizontalScrollPosition = scrollArea->horizontalScrollBar()->value();
			break;
		}
		ancestor = ancestor->parentWidget();
	}

	HideResultsWidgets(true);
	m_status->setText("Running...");
	m_resultsTableWidget->clearContents();
	m_preferredBase->setText("Not available");
	m_confidence->setText("Not available");
	m_startButton->setHidden(true);
	m_worker = new BaseAddressDetectionThread(&m_inputs, m_view);
	connect(m_worker, &BaseAddressDetectionThread::ResultReady, this, &BaseAddressDetectionWidget::HandleResults);
	connect(m_worker, &BaseAddressDetectionThread::finished, m_worker, &QObject::deleteLater);
	m_worker->start();
	m_abortButton->setHidden(false);

	// Restore scroll position after widget visibility changes
	if (scrollArea)
	{
		scrollArea->verticalScrollBar()->setValue(savedVerticalScrollPosition);
		scrollArea->horizontalScrollBar()->setValue(savedHorizontalScrollPosition);
	}
}


void BaseAddressDetectionWidget::Abort()
{
	m_worker->Abort();
	m_abortButton->setHidden(true);
	m_startButton->setHidden(false);
	m_startButton->setEnabled(false);
}


const std::string BaseAddressDetectionWidget::GetRebaseViewName()
{
	auto fileMetadata = m_view->GetFile();
	if (!fileMetadata)
		return "";

	for (const auto& viewName : fileMetadata->GetExistingViews())
	{
		if (viewName != "Raw")
			return viewName;
	}

	return "";
}


void BaseAddressDetectionWidget::RebaseWithFullAnalysis()
{
	auto fileMetadata = m_view->GetFile();
	if (!fileMetadata)
		return;

	uint64_t address;
	string errorStr;
	if (!BinaryNinja::BinaryView::ParseExpression(m_view, m_reloadBase->text().toStdString(), address, 0, errorStr))
	{
		m_status->setText(QString("Invalid rebase address (%1)").arg(QString::fromStdString(errorStr)));
		return;
	}

	auto rebaseViewName = GetRebaseViewName();
	if (!rebaseViewName.empty())
	{
		// Found an existing view that isn't raw, rebase it
		auto view = m_view->GetFile()->GetViewOfType(rebaseViewName);
		if (!view)
			return;

		bool result = false;
		ProgressTask* task = new ProgressTask(this, "Rebase", "Rebasing...", "Cancel", [&](BinaryNinja::ProgressFunction progress) {
			result = fileMetadata->Rebase(view, address, progress);
		});
		task->wait();
		if (!result)
			return;

		view->Reanalyze();
	}
	else
	{
		// Only a raw view exists - load the binary and run full analysis
		BinaryNinja::Settings::Instance()->Set("analysis.mode", "full", m_view);
		map<string, BinaryNinja::Ref<BinaryNinja::Metadata>> metadataMap = {
			{"analysis.linearSweep.permissive", new BinaryNinja::Metadata(true)},
			{"loader.imageBase", new BinaryNinja::Metadata((uint64_t) address)},
		};

		if (m_inputs.ArchitectureBox->currentText() != "auto detect")
			metadataMap["loader.platform"] = new BinaryNinja::Metadata(m_inputs.ArchitectureBox->currentText().toStdString());

		auto options = new BinaryNinja::Metadata(metadataMap);
		auto newView = Load(m_view->GetFile()->GetViewOfType("Raw"), false, options->GetJsonString());
		if (!newView)
			return;

		rebaseViewName = newView->GetTypeName();
	}

	// Refresh the UI and jump to Linear view
	auto frame = ViewFrame::viewFrameForWidget(this);
	if (!frame)
		return;

	auto fileContext = frame->getFileContext();
	if (!fileContext)
		return;

	auto uiContext = UIContext::contextForWidget(this);
	if (!uiContext)
		return;

	fileContext->refreshDataViewCache();
	uiContext->recreateViewFrames(fileContext);
	QCoreApplication::processEvents();

	auto newFrame = fileContext->getCurrentViewFrame();
	if (!newFrame)
		return;

	newFrame->navigate(QString("Linear:%1").arg(QString::fromStdString(rebaseViewName)), address);
}


void BaseAddressDetectionWidget::CreateAdvancedSettingsGroup()
{
	int32_t row = 0;
	int32_t column = 0;
	auto grid = new QGridLayout();

	grid->addWidget(new QLabel("Min. String Length:"), row, column, Qt::AlignLeft);
	m_inputs.StrlenLineEdit = new QLineEdit("0n10");
	grid->addWidget(m_inputs.StrlenLineEdit, row, column + 1, Qt::AlignLeft);

	grid->addWidget(new QLabel("Alignment:"), row, column + 2, Qt::AlignLeft);
	m_inputs.AlignmentLineEdit = new QLineEdit("0n1024");
	grid->addWidget(m_inputs.AlignmentLineEdit, row++, column + 3, Qt::AlignLeft);

	grid->addWidget(new QLabel("Lower Boundary:"), row, column, Qt::AlignLeft);
	m_inputs.LowerBoundary = new QLineEdit("0x0");
	grid->addWidget(m_inputs.LowerBoundary, row, column + 1, Qt::AlignLeft);

	grid->addWidget(new QLabel("Upper Boundary:"), row, column + 2, Qt::AlignLeft);
	m_inputs.UpperBoundary = new QLineEdit("0xffffffffffffffff");
	grid->addWidget(m_inputs.UpperBoundary, row++, column + 3, Qt::AlignLeft);

	grid->addWidget(new QLabel("Points Of Interest:"), row, column, Qt::AlignLeft);
	auto poiList = QStringList() << "All" << "Strings only" << "Functions only";
	m_inputs.POIBox = new QComboBox(this);
	m_inputs.POIBox->addItems(poiList);
	grid->addWidget(m_inputs.POIBox, row, column + 1, Qt::AlignLeft);

	grid->addWidget(new QLabel("Max Pointers:"), row, column + 2, Qt::AlignLeft);
	m_inputs.MaxPointersPerCluster = new QLineEdit("0n128");
	grid->addWidget(m_inputs.MaxPointersPerCluster, row++, column + 3, Qt::AlignLeft);

	m_advancedSettingsGroup = new ExpandableGroup(grid);
	m_advancedSettingsGroup->setTitle("Advanced Settings");
}


BaseAddressDetectionWidget::BaseAddressDetectionWidget(QWidget* parent,
	BinaryNinja::Ref<BinaryNinja::BinaryView> bv) : QWidget(parent)
{
	m_view = bv->GetParentView() ? bv->GetParentView() : bv;
	m_layout = new QGridLayout();
	int32_t row = 0;
	int32_t column = 0;

	m_layout->addWidget(new QLabel("Architecture:"), row, column, Qt::AlignLeft);
	m_inputs.ArchitectureBox = new QComboBox(this);
	auto architectures = BinaryNinja::Architecture::GetList();
	auto archItemList = QStringList();
	archItemList << "auto detect";
	for (const auto& arch : architectures)
		archItemList << QString::fromStdString(arch->GetName());
	m_inputs.ArchitectureBox->addItems(archItemList);
	m_layout->addWidget(m_inputs.ArchitectureBox, row++, column + 1, Qt::AlignLeft);

	m_layout->addWidget(new QLabel("Analysis Level:"), row, column, Qt::AlignLeft);
	m_inputs.AnalysisBox = new QComboBox(this);
	auto analysisItemList = QStringList() << "full" << "basic" << "controlFlow";
	m_inputs.AnalysisBox->addItems(analysisItemList);
	m_layout->addWidget(m_inputs.AnalysisBox, row++, column + 1, Qt::AlignLeft);

	CreateAdvancedSettingsGroup();
	m_layout->addWidget(m_advancedSettingsGroup, row++, column, 1, 4);

	m_startButton = new QPushButton("Start");
	connect(m_startButton, &QPushButton::clicked, this, &BaseAddressDetectionWidget::DetectBaseAddress);
	m_layout->addWidget(m_startButton, row, column, Qt::AlignLeft);

	m_abortButton = new QPushButton("Abort");
	connect(m_abortButton, &QPushButton::clicked, this, &BaseAddressDetectionWidget::Abort);
	m_abortButton->setHidden(true);
	m_layout->addWidget(m_abortButton, row, column, Qt::AlignLeft);

	m_status = new QLabel("Not running");
	auto palette = m_status->palette();
	palette.setColor(QPalette::WindowText, getThemeColor(AlphanumericHighlightColor));
	m_status->setPalette(palette);
	m_status->setFont(getMonospaceFont(this));
	m_layout->addWidget(m_status, row++, column + 1, 1, 2, Qt::AlignLeft);

	m_preferredBaseLabel = new QLabel("Preferred Base:");
	m_layout->addWidget(m_preferredBaseLabel, row, column, Qt::AlignLeft);
	m_preferredBase = new QLabel("Not available");
	m_preferredBase->setTextInteractionFlags(Qt::TextSelectableByMouse);
	m_preferredBase->setFont(getMonospaceFont(this));
	m_preferredBase->setPalette(palette);
	m_layout->addWidget(m_preferredBase, row, column + 1, Qt::AlignLeft);

	m_confidenceLabel = new QLabel("Confidence:");
	m_layout->addWidget(m_confidenceLabel, row, column + 2, Qt::AlignLeft);
	m_confidence = new QLabel("Not available");
	m_confidence->setFont(getMonospaceFont(this));
	m_confidence->setPalette(palette);
	m_layout->addWidget(m_confidence, row++, column + 3, Qt::AlignLeft);

	m_resultsTableWidget = new QTableWidget(this);
	m_resultsTableWidget->setColumnCount(5);
	QStringList header;
	header << "Base Address" << "Score" << "String Hits" << "Function Hits" << "Data Hits";
	m_resultsTableWidget->setHorizontalHeaderLabels(header);
	m_resultsTableWidget->horizontalHeader()->setDefaultAlignment(Qt::AlignLeft);
	m_resultsTableWidget->horizontalHeader()->setStretchLastSection(true);
	m_resultsTableWidget->verticalHeader()->setVisible(false);
	m_resultsTableWidget->setEditTriggers(QAbstractItemView::NoEditTriggers);
	m_resultsTableWidget->setSelectionBehavior(QAbstractItemView::SelectRows);
	m_resultsTableWidget->setSelectionMode(QAbstractItemView::SingleSelection);
	m_resultsTableWidget->setMinimumHeight(150);
	m_layout->addWidget(m_resultsTableWidget, row++, column, 1, 5);
	connect(m_resultsTableWidget, &QTableWidget::clicked, this, &BaseAddressDetectionWidget::GetClickedBaseAddress);

	m_reloadBase = new QLineEdit("0x0");
	m_layout->addWidget(m_reloadBase, row, column, Qt::AlignLeft);

	m_rebaseButton = new QPushButton("Start Full Analysis");
	connect(m_rebaseButton, &QPushButton::clicked, this, &BaseAddressDetectionWidget::RebaseWithFullAnalysis);
	m_layout->addWidget(m_rebaseButton, row, column + 1, Qt::AlignLeft);

	HideResultsWidgets(true);
	m_layout->setColumnStretch(3, 1);
	setLayout(m_layout);
}