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

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Net.Http.Json;
using System.Net.Mime;
using System.Threading.Tasks;
using System.Xml;
using AutomationTool.Tasks.CloudDDC;
using EpicGames.Core;
using EpicGames.Horde.Artifacts;
using EpicGames.Horde.Commits;
using EpicGames.Horde.Streams;
using EpicGames.Serialization;
using Microsoft.Extensions.Logging;
using UnrealBuildBase;
using UnrealBuildTool;

#nullable enable

namespace AutomationTool.Tasks
{
	/// <summary>
	/// Parameters for a <see cref="RetrieveCloudArtifactTask"/>.
	/// </summary>
	public class RetrieveCloudArtifactTaskParameters
	{
		/// <summary>
		/// Stream containing the artifact
		/// </summary>
		[TaskParameter(Optional = true)]
		public string? StreamId { get; set; } = null!;

		/// <summary>
		/// Change number for the artifact
		/// </summary>
		[TaskParameter(Optional = true)]
		public string? Commit { get; set; }

		/// <summary>
		/// Maximum commit for the artifact
		/// </summary>
		[TaskParameter(Optional = true)]
		public string? MaxCommit { get; set; }

		/// <summary>
		/// Name of the artifact
		/// </summary>
		[TaskParameter(Optional = true)]
		public string Name { get; set; } = null!;

		/// <summary>
		/// The artifact type. Determines the permissions and expiration policy for the artifact.
		/// </summary>
		[TaskParameter]
		public string Type { get; set; } = null!;

		/// <summary>
		/// Keys for the artifact
		/// </summary>
		[TaskParameter(Optional = true)]
		public string Keys { get; set; } = null!;

		/// <summary>
		/// Output directory for 
		/// </summary>
		[TaskParameter]
		public string? OutputDir { get; set; }

		/// <summary>
		/// The platform the artifact is created for
		/// </summary>
		[TaskParameter]
		public string Platform { get; set; } = null!;

		/// <summary>
		/// The path to the uproject this artifact is created for
		/// </summary>
		[TaskParameter]
		public FileReference Project { get; set; } = null!;

		/// <summary>
		/// The platform the artifact is created for
		/// </summary>
		[TaskParameter(Optional = true)]
		public string? Host { get; set; }

		/// <summary>
		/// The platform the artifact is created for
		/// </summary>
		[TaskParameter(Optional = true)]
		public string? Namespace { get; set; }

		/// <summary>
		/// The access token to use
		/// </summary>
		[TaskParameter(Optional = true)]
		public string? AccessToken { get; set; }

		/// <summary>
		/// Set this to use the latest match if multiple artifacts are possible matches
		/// </summary>
		[TaskParameter(Optional = true)]
		public bool AllowMultipleMatches { get; set; } = false;

		/// <summary>
		/// Enable to use multipart endpoints if valuable
		/// </summary>
		[TaskParameter(Optional = true)]
		public bool AllowMultipart { get; set; } = false;

		/// <summary>
		/// Set the explicit http version to use. None to use http handshaking.
		/// </summary>
		[TaskParameter(Optional = true)]
		public string HttpVersion { get; set; } = "None";

		/// <summary>
		/// Increase the number of worker threads used by zen, may cause machine to be less responsive but will generally improve download times
		/// </summary>
		[TaskParameter(Optional = true)]
		public bool BoostWorkers { get; set; } = false;

		/// <summary>
		/// Enable to create an Unreal Insights trace of the download process
		/// </summary>
		public bool EnableTracing { get; set; } = true;
	}

	/// <summary>
	/// Retrieves an artifact from Cloud DDC
	/// </summary>
	[TaskElement("RetrieveCloudArtifact", typeof(RetrieveCloudArtifactTaskParameters))]
	public class RetrieveCloudArtifactTask : BgTaskImpl
	{
		readonly RetrieveCloudArtifactTaskParameters _parameters;

		/// <summary>
		/// Constructor.
		/// </summary>
		/// <param name="parameters">Parameters for this task.</param>
		public RetrieveCloudArtifactTask(RetrieveCloudArtifactTaskParameters parameters)
			=> _parameters = parameters;

		/// <summary>
		/// ExecuteAsync the task.
		/// </summary>
		/// <param name="job">Information about the current job.</param>
		/// <param name="buildProducts">Set of build products produced by this node.</param>
		/// <param name="tagNameToFileSet">Mapping from tag names to the set of files they include.</param>
		public override async Task ExecuteAsync(JobContext job, HashSet<FileReference> buildProducts, Dictionary<string, HashSet<FileReference>> tagNameToFileSet)
		{
			if (!String.IsNullOrEmpty(_parameters.Commit) && !String.IsNullOrEmpty(_parameters.MaxCommit))
			{
				throw new AutomationException("Cannot specify both Commit and MaxCommit parameters for retrieving an artifact.");
			}

			string project = _parameters.Project.GetFileNameWithoutExtension();
			FileReference? projectPath = _parameters.Project;
			if (!File.Exists(_parameters.Project.FullName))
			{
				// if the project is not an uproject that exists we do the scanning for engine inis without the project specific overrides
				projectPath = null;
				Logger.LogInformation("\'{Project}\' is not a valid file. Loading engine configs without project overrides.", _parameters.Project);
			}
			string platform = _parameters.Platform;

			bool foundConfig = false;
			CloudConfiguration? cloudConfig = null;

			if (projectPath != null)
			{
				(Dictionary<UnrealTargetPlatform, ConfigHierarchy> engineConfigs, Dictionary<UnrealTargetPlatform, ConfigHierarchy> _) = CreateCloudArtifactTask.GetIniConfigs(projectPath);
				ConfigHierarchy config = engineConfigs[HostPlatform.Current.HostEditorPlatform];
				foundConfig = config.TryGetValueGeneric("StorageServers", "Cloud", out CloudConfiguration foundCloudConfig);
				cloudConfig = foundCloudConfig;
			}

			// Figure out the current change and stream id
			StreamId streamId;
			if (!String.IsNullOrEmpty(_parameters.StreamId))
			{
				streamId = new StreamId(_parameters.StreamId);
			}
			else
			{
				string? streamIdEnvVar = Environment.GetEnvironmentVariable("UE_HORDE_STREAMID");
				if (!String.IsNullOrEmpty(streamIdEnvVar))
				{
					streamId = new StreamId(streamIdEnvVar);
				}
				else
				{
					throw new AutomationException("Missing UE_HORDE_STREAMID environment variable; unable to determine current stream.");
				}
			}

			string? cloudHostEnvVar = Environment.GetEnvironmentVariable(OperatingSystem.IsWindows() ? "UE-CloudPublishHost" : "UE_CloudPublishHost");
			
			string host;
			if (!String.IsNullOrEmpty(_parameters.Host))
			{
				host = _parameters.Host;
			}
			else if (!String.IsNullOrEmpty(cloudHostEnvVar))
			{
				host = cloudHostEnvVar;
			}
			else if (foundConfig && cloudConfig != null && !String.IsNullOrEmpty(cloudConfig.Value.Host))
			{
				string cloudConfigHost = cloudConfig.Value.Host;
				if (cloudConfigHost.Contains(';', StringComparison.OrdinalIgnoreCase))
				{
					// if its a list pick the first element
					cloudConfigHost = cloudConfigHost.Split(";").First();
				}

				host = cloudConfigHost;
			}
			else
			{
				throw new AutomationException("Missing UE-CloudPublishHost environment variable; unable to determine cloud host. Specify the environment variable or define it in your StorageServer section in your engine ini");
			}

			string ns;
			string? cloudDefaultNamespaceEnv = Environment.GetEnvironmentVariable(OperatingSystem.IsWindows() ? "UE-CloudPublishNamespace" : "UE_CloudPublishNamespace");
				
			if (!String.IsNullOrEmpty(_parameters.Namespace))
			{
				ns = _parameters.Namespace;
			}
			else if (!String.IsNullOrEmpty(cloudDefaultNamespaceEnv))
			{
				ns = cloudDefaultNamespaceEnv;
			}
			else if (foundConfig && cloudConfig != null && !String.IsNullOrEmpty(cloudConfig.Value.BuildsNamespace))
			{
				ns = cloudConfig.Value.BuildsNamespace;
			}
			else
			{
				throw new AutomationException("Missing UE-CloudPublishNamespace environment variable; unable to default namespace please specify it in the task or define it in your StorageServer section in your engine ini.");
			}
			string httpVersion = _parameters.HttpVersion;
			if (String.IsNullOrEmpty(httpVersion) || httpVersion == "None")
			{
				string? httpVersionEnvironmentVariable = Environment.GetEnvironmentVariable(OperatingSystem.IsWindows() ? "UE-CloudPublishHttpVersion": "UE_CloudPublishHttpVersion");
				if (!String.IsNullOrEmpty(httpVersionEnvironmentVariable))
				{
					httpVersion = httpVersionEnvironmentVariable;
				}
			}

			string accessToken;
			if (!String.IsNullOrEmpty(_parameters.AccessToken))
			{
				accessToken = _parameters.AccessToken;
			}
			else
			{
				string? cloudAccessTokenEnvVar = Environment.GetEnvironmentVariable(OperatingSystem.IsWindows() ? "UE-CloudDataCacheAccessToken" : "UE_CloudDataCacheAccessToken");
				if (!String.IsNullOrEmpty(cloudAccessTokenEnvVar))
				{
					accessToken = cloudAccessTokenEnvVar;
				}
				else
				{
					throw new AutomationException("Missing UE-CloudDataCacheAccessToken environment variable; unable to find access token to use.");
				}
			}

			// Get the current commit id
			CommitId? minCommitId = null;
			CommitId? maxCommitId = null;
			if (!String.IsNullOrEmpty(_parameters.MaxCommit))
			{
				maxCommitId = new CommitId(_parameters.MaxCommit);
			}
			else if (!String.IsNullOrEmpty(_parameters.Commit))
			{
				minCommitId = maxCommitId = new CommitId(_parameters.Commit);
			}
			else
			{
				try
				{
					int change = CommandUtils.P4Env.Changelist;
					if (change > 0)
					{
						minCommitId = maxCommitId = CommitId.FromPerforceChange(CommandUtils.P4Env.Changelist);
					}
				}
				catch (AutomationException)
				{
					// not an error to run without p4
				}
			}

			string? name = null;
			if (!String.IsNullOrEmpty(_parameters.Name))
			{
				name = _parameters.Name;
			}

			ArtifactType type = new ArtifactType(_parameters.Type);
			List<string> keys = (_parameters.Keys ?? String.Empty).Split(';', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries).ToList();

			CbWriter queryWriter = new CbWriter();
			queryWriter.BeginObject(); // root
			queryWriter.BeginObject("query"); // query
			
			queryWriter.BeginObject("stream");
			queryWriter.WriteString("$eq", CreateCloudArtifactTask.SanitizeBucketValue(streamId.ToString()));
			queryWriter.EndObject();

			if (maxCommitId != null)
			{
				if (minCommitId == maxCommitId)
				{
					queryWriter.BeginObject("commit");
					queryWriter.WriteString("$eq", maxCommitId.ToString());
					queryWriter.EndObject();
				}
				else if (minCommitId != null)
				{
					queryWriter.BeginObject("commit");
					queryWriter.WriteString("$lte", maxCommitId.ToString());
					queryWriter.WriteString("$gte", minCommitId.ToString());
					queryWriter.EndObject();
				}
				else
				{
					queryWriter.BeginObject("commit");
					queryWriter.WriteString("$lte", maxCommitId.ToString());
					queryWriter.EndObject();
				}
			}

			if (name != null)
			{
				queryWriter.BeginObject("name");
				queryWriter.WriteString("$eq", name);
				queryWriter.EndObject();
			}

			queryWriter.BeginObject("type");
			queryWriter.WriteString("$eq", type.ToString());
			queryWriter.EndObject();

			if (keys.Count > 0)
			{
				queryWriter.BeginObject("keys");
				queryWriter.BeginUniformArray("$in", CbFieldType.String);
				foreach (string key in keys)
				{
					queryWriter.WriteStringValue(key);
				}
				queryWriter.EndUniformArray();
				queryWriter.EndObject();
			}
			
			queryWriter.EndObject(); // end query
			queryWriter.BeginObject("options");
			queryWriter.WriteInteger("limit", 1); // we only care about the first match
			queryWriter.EndObject(); // options
			queryWriter.EndObject(); // end root object

			CbObject queryObject = queryWriter.ToObject();
			DirectoryReference outputDir = ResolveDirectory(_parameters.OutputDir);
			string bucketId = $"{CreateCloudArtifactTask.SanitizeBucketValue(project)}.{CreateCloudArtifactTask.SanitizeBucketValue(type.ToString())}.{CreateCloudArtifactTask.SanitizeBucketValue(streamId.ToString())}.{CreateCloudArtifactTask.SanitizeBucketValue(platform)}".ToLowerInvariant();
			using HttpClient httpClient = BuildHttpClient(host, ns, bucketId, accessToken);

			List<CbObjectId> artifactIds = await SearchArtifactsAsync(httpClient, queryObject);

			if (artifactIds.Count == 0)
			{
				throw new AutomationException($"Unable to find any artifact matching given criteria in namespace {ns} and bucket {bucketId} criteria: {queryObject.ToJson()}");
			}

			if (!_parameters.AllowMultipleMatches && artifactIds.Count != 1)
			{
				throw new AutomationException("More then one matching artifact given criteria, set \"AllowMultipleMatches\" option if you want to use the newest matching artifact or refine the search");
			}
			
			CbObjectId artifact = artifactIds.First();
			Logger.LogInformation("Found artifact {ArtifactId}", artifact);

			FileReference? traceFile = null;
			if (_parameters.EnableTracing)
			{
				// this is automatically uploaded as it's in the Saved folder
				traceFile = FileReference.Combine(Unreal.RootDirectory, "Engine/Programs/AutomationTool/Saved/Logs/build-download.utrace");
			}

			BuildDownload.DownloadBuild(host, outputDir, ns, bucketId, artifact.ToString(), accessToken, allowMultipart: _parameters.AllowMultipart, assumeHttp2: String.Equals(httpVersion, "http2-only", StringComparison.OrdinalIgnoreCase), boostWorkers: _parameters.BoostWorkers, traceFile: traceFile);
		}

		static HttpClient BuildHttpClient(string host, string ns, string bucketId, string accessToken)
		{
			HttpClient httpClient = new HttpClient();
			string url = $"{host}/api/v2/builds/{ns}/{bucketId}/";
			try
			{
				httpClient.BaseAddress = new Uri(url);
			}
			catch (UriFormatException)
			{
				throw new AutomationException($"{url} is not a valid url. Make sure host, namespace and bucket have been specified");
			}

			CbObjectId sessionId = CbObjectId.NewObjectId();
			Logger.LogInformation("Using SessionId {SessionId} please include this in any error reports.", sessionId);
			string authScheme = "Bearer"; // TODO: for a proper implementation this should be configurable
			httpClient.DefaultRequestHeaders.Add("Authorization", $"{authScheme} {accessToken}");
			httpClient.DefaultRequestHeaders.Add("UE-Session", sessionId.ToString());
			httpClient.DefaultRequestHeaders.Add("Accept", CustomMediaTypeNames.UnrealCompactBinary);
			httpClient.Timeout = TimeSpan.FromMinutes(5.0); // bump timeout for each request as we sometimes have larger files
			return httpClient;
		}
		static async Task<List<CbObjectId>> SearchArtifactsAsync(HttpClient httpClient, CbObject queryObject)
		{
			using StringContent content = new StringContent(queryObject.ToJson());
			content.Headers.ContentType = new MediaTypeHeaderValue(MediaTypeNames.Application.Json);

			using HttpRequestMessage request = new HttpRequestMessage(HttpMethod.Post, new Uri($"search", UriKind.Relative));
			request.Content = content;
			request.Headers.Add("Accept", MediaTypeNames.Application.Json);
			using HttpResponseMessage result = await httpClient.SendAsync(request);
			result.EnsureSuccessStatusCode();

			SearchResult? searchResult = await result.Content.ReadFromJsonAsync<SearchResult>();
			return searchResult!.Results.Select(r => r.BuildId).ToList();
		}

		/// <inheritdoc/>
		public override void Write(XmlWriter writer)
			=> Write(writer, _parameters);

		/// <inheritdoc/>
		public override IEnumerable<string> FindConsumedTagNames()
			=> Enumerable.Empty<string>();

		/// <inheritdoc/>
		public override IEnumerable<string> FindProducedTagNames()
			=> Enumerable.Empty<string>();
	}

	class BuildSearchResult
	{
		[CbField("buildId")]
		public CbObjectId BuildId { get; set; }
	}

	class SearchResult
	{
		[CbField("results")] public List<BuildSearchResult> Results { get; set; } = new List<BuildSearchResult>();

		[CbField("partialResult")] public bool PartialResult { get; set; }

		public SearchResult()
		{
			Results = new List<BuildSearchResult>();
			PartialResult = false;
		}
	}

	static class BuildDownload
	{
		public static void DownloadBuild(string host, DirectoryReference targetDir, string namespaceId, string bucketId, string buildId, string accessToken, bool allowMultipart = false, bool assumeHttp2 = false, bool boostWorkers = false, FileReference? traceFile = null)
		{
			FileInfo zenExe = new FileInfo("Engine/Binaries/Win64/zen.exe");

			string http2Options = assumeHttp2 ? "--assume-http2" : "";
			string boostWorkersOption = boostWorkers ? "--boost-workers" : "";
			// trace file option needs to be passed before other arguments including the verbs
			string traceOption = traceFile != null ? $"--tracefile={traceFile.FullName} " : "";
			// pass the access token via environment variable to avoid showing it on the cli
			string cmdline =
				$"{traceOption}builds download --url {host} --namespace {namespaceId} --bucket {bucketId} --access-token-env UE-CloudDataCacheAccessToken --plain-progress --allow-multipart={allowMultipart} {http2Options} {boostWorkersOption} \"{targetDir}\" {buildId}";
			CommandUtils.RunAndLog(CommandUtils.CmdEnv, zenExe.FullName, cmdline, Options: CommandUtils.ERunOptions.Default,
				EnvVars: new Dictionary<string, string> { { "UE-CloudDataCacheAccessToken", accessToken } });
		}
	}
}