// Copyright Epic Games, Inc. All Rights Reserved.

#include <zencompute/mockimds.h>

#include <zencore/fmtutils.h>

#if ZEN_WITH_TESTS

namespace zen::compute {

const char*
MockImdsService::BaseUri() const
{
	return "/";
}

void
MockImdsService::HandleRequest(HttpServerRequest& Request)
{
	std::string_view Uri = Request.RelativeUri();

	// AWS endpoints live under /latest/
	if (Uri.starts_with("latest/"))
	{
		if (ActiveProvider == CloudProvider::AWS)
		{
			HandleAwsRequest(Request);
			return;
		}
		Request.WriteResponse(HttpResponseCode::NotFound);
		return;
	}

	// Azure endpoints live under /metadata/
	if (Uri.starts_with("metadata/"))
	{
		if (ActiveProvider == CloudProvider::Azure)
		{
			HandleAzureRequest(Request);
			return;
		}
		Request.WriteResponse(HttpResponseCode::NotFound);
		return;
	}

	// GCP endpoints live under /computeMetadata/
	if (Uri.starts_with("computeMetadata/"))
	{
		if (ActiveProvider == CloudProvider::GCP)
		{
			HandleGcpRequest(Request);
			return;
		}
		Request.WriteResponse(HttpResponseCode::NotFound);
		return;
	}

	Request.WriteResponse(HttpResponseCode::NotFound);
}

// ---------------------------------------------------------------------------
// AWS
// ---------------------------------------------------------------------------

void
MockImdsService::HandleAwsRequest(HttpServerRequest& Request)
{
	std::string_view Uri = Request.RelativeUri();

	// IMDSv2 token acquisition (PUT only)
	if (Uri == "latest/api/token" && Request.RequestVerb() == HttpVerb::kPut)
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Aws.Token);
		return;
	}

	// Instance identity
	if (Uri == "latest/meta-data/instance-id")
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Aws.InstanceId);
		return;
	}

	if (Uri == "latest/meta-data/placement/availability-zone")
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Aws.AvailabilityZone);
		return;
	}

	if (Uri == "latest/meta-data/instance-life-cycle")
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Aws.LifeCycle);
		return;
	}

	// Autoscaling lifecycle state — 404 when not in an ASG
	if (Uri == "latest/meta-data/autoscaling/target-lifecycle-state")
	{
		if (Aws.AutoscalingState.empty())
		{
			Request.WriteResponse(HttpResponseCode::NotFound);
			return;
		}
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Aws.AutoscalingState);
		return;
	}

	// Spot interruption notice — 404 when no interruption pending
	if (Uri == "latest/meta-data/spot/instance-action")
	{
		if (Aws.SpotAction.empty())
		{
			Request.WriteResponse(HttpResponseCode::NotFound);
			return;
		}
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Aws.SpotAction);
		return;
	}

	Request.WriteResponse(HttpResponseCode::NotFound);
}

// ---------------------------------------------------------------------------
// Azure
// ---------------------------------------------------------------------------

void
MockImdsService::HandleAzureRequest(HttpServerRequest& Request)
{
	std::string_view Uri = Request.RelativeUri();

	// Instance metadata (single JSON document)
	if (Uri == "metadata/instance")
	{
		std::string Json = fmt::format(R"({{"compute":{{"vmId":"{}","location":"{}","priority":"{}","vmScaleSetName":"{}"}}}})",
									   Azure.VmId,
									   Azure.Location,
									   Azure.Priority,
									   Azure.VmScaleSetName);

		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Json);
		return;
	}

	// Scheduled events for termination monitoring
	if (Uri == "metadata/scheduledevents")
	{
		std::string Json;
		if (Azure.ScheduledEventType.empty())
		{
			Json = R"({"Events":[]})";
		}
		else
		{
			Json = fmt::format(R"({{"Events":[{{"EventType":"{}","EventStatus":"{}"}}]}})",
							   Azure.ScheduledEventType,
							   Azure.ScheduledEventStatus);
		}

		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Json);
		return;
	}

	Request.WriteResponse(HttpResponseCode::NotFound);
}

// ---------------------------------------------------------------------------
// GCP
// ---------------------------------------------------------------------------

void
MockImdsService::HandleGcpRequest(HttpServerRequest& Request)
{
	std::string_view Uri = Request.RelativeUri();

	if (Uri == "computeMetadata/v1/instance/id")
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Gcp.InstanceId);
		return;
	}

	if (Uri == "computeMetadata/v1/instance/zone")
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Gcp.Zone);
		return;
	}

	if (Uri == "computeMetadata/v1/instance/scheduling/preemptible")
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Gcp.Preemptible);
		return;
	}

	if (Uri == "computeMetadata/v1/instance/maintenance-event")
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kText, Gcp.MaintenanceEvent);
		return;
	}

	Request.WriteResponse(HttpResponseCode::NotFound);
}

}  // namespace zen::compute

#endif	// ZEN_WITH_TESTS