using System;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text;
using System.Threading;
using Tesla;

namespace VPod;

/// <summary>
/// The pod side of the SecureConfig first-boot provisioning protocol, display-only:
/// unlike a real pod (the Launcher's PodSecureConfigurator) it never touches the
/// NIC, registry or hostname — the network config the console assigns is only
/// surfaced to the UI. The wire behaviour matches the real pod, so the console's
/// Manage Site "Configure" flow works unmodified:
///
///   1. Broadcast a "RQST" beacon (MAC + 3-char RequestId) to UDP 53291 every 10 s
///      — the console shows a "Configure &lt;RequestId&gt;" button.
///   2. Operator enters the pod's network settings and the 5-char passphrase shown
///      in vPOD's window; the console broadcasts an AES-encrypted "RPLY" (network
///      config) to UDP 53292, key = PBKDF2(passphrase).
///   3. The console TCP-connects to the pod's (entered) address on 53292; after the
///      OFB/CONF handshake on the passphrase key, the pod sends an RSA public key
///      and receives the RSA-encrypted 32-byte session key — the key that unlocks
///      the launcher RPC channel (TCP 53290) from then on.
///
/// Reuses the shared TeslaSecureConfiguration pieces where they are public
/// (UdpBeacon, BasicConfigResponse, NegotiateCryptoStreams); the passphrase KDF
/// salt and the "RQST"/"RPLY" tags are internal there and duplicated below.
/// </summary>
internal sealed class PodProvisioning
{
	public const int ConsoleRequestPort = 53291; // console listens for RQST beacons
	public const int PodReplyPort = 53292;       // pod listens for RPLY + the TCP key exchange

	// Mirrors of internals in SecureConfig/SecureConfig.cs (PodConfigurationServer):
	// the PBKDF2 salt for the passphrase-derived AES key, and the pod-side
	// passphrase/request-id alphabet + lengths (SetupPod validates passphrase == 5).
	private static readonly byte[] sPassphraseSalt = new byte[32]
	{
		23, 171, 81, 217, 236, 209, 212, 116, 169, 9,
		74, 52, 39, 251, 31, 242, 222, 196, 249, 241,
		166, 216, 158, 218, 21, 17, 71, 101, 50, 231,
		231, 239
	};
	private const string Alphabet = "23456789ABCDEFGHJKLMNPQRSTUVWXYZ";
	private const int RequestIdLength = 3;
	private const int PassphraseLength = 5;

	private readonly byte[] mMacAddress;
	private readonly object mLock = new object();
	private UdpBeacon mBeacon;
	private UdpClient mReplyListener;
	private TcpListener mKeyExchangeListener;
	private Thread mWorker;
	private volatile bool mRunning;
	// Bumped by every Start/Stop so a worker from a previous session can neither
	// tear down nor complete a newer one (e.g. quick power-off/power-on cycles).
	private volatile int mGeneration;

	public event Action<string> Log;
	public event Action<BasicConfigResponse> ConfigReceived; // display-only network config
	public event Action<byte[]> Provisioned;                 // the 32-byte session key

	public string RequestId { get; private set; }
	public string Passphrase { get; private set; }
	public bool IsRunning => mRunning;

	/// <summary>The pod's stable fake MAC: locally-administered "VPOD" + host id, so
	/// the console recognizes the same virtual pod across reprovisions (Site.FindPod).</summary>
	public static byte[] MacForHost(int hostId)
	{
		return new byte[6] { 0x02, 0x56, 0x50, 0x4F, 0x44, (byte)hostId };
	}

	public PodProvisioning(byte[] macAddress)
	{
		mMacAddress = macAddress;
	}

	public void Start()
	{
		int generation;
		lock (mLock)
		{
			if (mRunning)
			{
				return;
			}
			mRunning = true;
			generation = ++mGeneration;
			RequestId = GenerateRandomString(RequestIdLength);
			Passphrase = GenerateRandomString(PassphraseLength);

			byte[] payload = new byte[mMacAddress.Length + RequestIdLength];
			mMacAddress.CopyTo(payload, 0);
			Encoding.ASCII.GetBytes(RequestId).CopyTo(payload, mMacAddress.Length);
			mBeacon = new UdpBeacon(Encoding.UTF8.GetBytes("RQST"), payload, 10000.0, ConsoleRequestPort, null);
			mBeacon.Start();

			mWorker = new Thread(() => ProvisionWorker(generation)) { IsBackground = true, Name = "vPOD-provision" };
			mWorker.Start();
		}
		Log?.Invoke($"Provisioning: beaconing RQST (Request ID {RequestId}, passphrase {Passphrase}).");
	}

	public void Stop()
	{
		lock (mLock)
		{
			mGeneration++; // orphan any live worker
			if (!mRunning)
			{
				return;
			}
			mRunning = false;
			try { mBeacon?.Stop(); } catch { }
			mBeacon = null;
			try { mReplyListener?.Close(); } catch { }
			mReplyListener = null;
			try { mKeyExchangeListener?.Stop(); } catch { }
			mKeyExchangeListener = null;
		}
	}

	private bool IsCurrent(int generation)
	{
		return mRunning && generation == mGeneration;
	}

	private void ProvisionWorker(int generation)
	{
		try
		{
			// ---- Phase 1: wait for the console's RPLY (proves the operator typed
			// our passphrase) and surface the assigned network config. ----
			byte[] weakKey = DeriveKeyFromPassphrase(Passphrase);
			BasicConfigResponse config = ReceiveReply(weakKey, generation);
			if (config == null || !IsCurrent(generation))
			{
				return; // stopped
			}
			Log?.Invoke($"Provisioning: RPLY received — assigned IP {config.Address} / {config.Mask}" +
				(string.IsNullOrEmpty(config.HostName) ? "" : $", host \"{config.HostName}\"") +
				" (display only, not applied).");
			ConfigReceived?.Invoke(config);
			lock (mLock)
			{
				try { mBeacon?.Stop(); } catch { }
				mBeacon = null;
			}

			// ---- Phase 2: accept the console's TCP key exchange on 53292. ----
			TcpListener listener;
			lock (mLock)
			{
				if (!IsCurrent(generation))
				{
					return;
				}
				try
				{
					listener = new TcpListener(IPAddress.Any, PodReplyPort);
					listener.Start();
				}
				catch (Exception ex)
				{
					Log?.Invoke($"Provisioning: cannot listen on TCP {PodReplyPort}: {ex.Message}");
					return;
				}
				mKeyExchangeListener = listener;
			}
			Log?.Invoke("Provisioning: waiting for the console's key exchange on TCP " + PodReplyPort + "...");
			while (IsCurrent(generation))
			{
				TcpClient client = null;
				try
				{
					client = listener.AcceptTcpClient();
					byte[] sessionKey = ExchangeSessionKey(client, weakKey);
					if (sessionKey == null)
					{
						Log?.Invoke("Provisioning: key-exchange handshake failed (wrong passphrase key?); still waiting.");
						continue;
					}
					if (!IsCurrent(generation))
					{
						return;
					}
					Log?.Invoke("Provisioning: session key received — pod is provisioned.");
					Provisioned?.Invoke(sessionKey);
					return;
				}
				catch (Exception ex)
				{
					if (IsCurrent(generation))
					{
						Log?.Invoke("Provisioning: key exchange error: " + ex.Message);
					}
					else
					{
						return; // listener stopped
					}
				}
				finally
				{
					try { client?.Close(); } catch { }
				}
			}
		}
		finally
		{
			StopGeneration(generation);
		}
	}

	/// <summary>Tears the session down only if it is still the one this worker
	/// belongs to — a newer Start() must not be disturbed by an old worker exiting.</summary>
	private void StopGeneration(int generation)
	{
		lock (mLock)
		{
			if (generation != mGeneration)
			{
				return;
			}
			mRunning = false;
			try { mBeacon?.Stop(); } catch { }
			mBeacon = null;
			try { mReplyListener?.Close(); } catch { }
			mReplyListener = null;
			try { mKeyExchangeListener?.Stop(); } catch { }
			mKeyExchangeListener = null;
		}
	}

	/// <summary>Listens on UDP 53292 for an "RPLY" datagram that decrypts and
	/// verifies under our passphrase key. Cancellable mirror of the shared
	/// UdpBeaconListener (which cannot be stopped once blocked in Receive):
	/// packet = "RPLY"(4) + IV(16) + AES-CBC(config + SHA1(config)).</summary>
	private BasicConfigResponse ReceiveReply(byte[] weakKey, int generation)
	{
		UdpClient udp;
		lock (mLock)
		{
			if (!IsCurrent(generation))
			{
				return null;
			}
			try
			{
				udp = new UdpClient(PodReplyPort) { EnableBroadcast = true };
			}
			catch (Exception ex)
			{
				Log?.Invoke($"Provisioning: cannot listen on UDP {PodReplyPort}: {ex.Message}");
				return null;
			}
			mReplyListener = udp;
		}
		byte[] header = Encoding.UTF8.GetBytes("RPLY");
		using (Rijndael aes = Rijndael.Create())
		using (SHA1 sha1 = SHA1.Create())
		{
			aes.Key = weakKey;
			while (IsCurrent(generation))
			{
				byte[] packet;
				try
				{
					IPEndPoint remote = new IPEndPoint(IPAddress.Any, 0);
					packet = udp.Receive(ref remote);
				}
				catch
				{
					return null; // socket closed by Stop()
				}
				try
				{
					int ivLength = aes.IV.Length;
					if (packet.Length < header.Length + ivLength)
					{
						continue;
					}
					bool tagOk = true;
					for (int i = 0; i < header.Length; i++)
					{
						if (packet[i] != header[i]) { tagOk = false; break; }
					}
					if (!tagOk)
					{
						continue;
					}
					byte[] iv = new byte[ivLength];
					Buffer.BlockCopy(packet, header.Length, iv, 0, ivLength);
					byte[] plain;
					using (ICryptoTransform decryptor = aes.CreateDecryptor(aes.Key, iv))
					using (MemoryStream ms = new MemoryStream())
					{
						using (CryptoStream cs = new CryptoStream(ms, decryptor, CryptoStreamMode.Write))
						{
							cs.Write(packet, header.Length + ivLength, packet.Length - header.Length - ivLength);
							cs.FlushFinalBlock();
						}
						plain = ms.ToArray();
					}
					int messageLength = plain.Length - sha1.HashSize / 8;
					if (messageLength <= 0)
					{
						continue;
					}
					byte[] hash = sha1.ComputeHash(plain, 0, messageLength);
					bool hashOk = true;
					for (int i = 0; i < hash.Length; i++)
					{
						if (plain[messageLength + i] != hash[i]) { hashOk = false; break; }
					}
					if (!hashOk)
					{
						continue; // wrong passphrase (or noise) — keep listening
					}
					byte[] message = new byte[messageLength];
					Buffer.BlockCopy(plain, 0, message, 0, messageLength);
					return new BasicConfigResponse(message);
				}
				catch
				{
					// undecryptable/malformed datagram — keep listening
				}
			}
		}
		return null;
	}

	/// <summary>The RSA leg, mirroring the real pod (PodConfigurationClient): after
	/// the OFB/CONF handshake on the passphrase key, send our RSA public key and
	/// decrypt the console's session key with it. Returns null if CONF fails.</summary>
	private byte[] ExchangeSessionKey(TcpClient client, byte[] weakKey)
	{
		if (!PodConfigurationServer.NegotiateCryptoStreams(client.GetStream(), weakKey, out Stream outStream, out Stream inStream))
		{
			return null;
		}
		BinaryWriter writer = new BinaryWriter(outStream);
		BinaryReader reader = new BinaryReader(inStream);
		using (RSACryptoServiceProvider rsa = new RSACryptoServiceProvider(2048))
		{
			writer.Write(rsa.ToXmlString(includePrivateParameters: false));
			writer.Flush();
			byte[] encryptedKey = reader.ReadBytes(reader.ReadInt32());
			return rsa.Decrypt(encryptedKey, fOAEP: false);
		}
	}

	internal static byte[] DeriveKeyFromPassphrase(string passphrase)
	{
		using (Rfc2898DeriveBytes pbkdf2 = new Rfc2898DeriveBytes(passphrase, sPassphraseSalt, 1000))
		{
			return pbkdf2.GetBytes(32);
		}
	}

	private static string GenerateRandomString(int length)
	{
		// RNGCryptoServiceProvider-quality randomness is unnecessary here (the real
		// pod uses System.Random too), but avoid same-seed collisions across quick
		// restarts by seeding from Guid entropy.
		Random random = new Random(Guid.NewGuid().GetHashCode());
		char[] chars = new char[length];
		for (int i = 0; i < length; i++)
		{
			chars[i] = Alphabet[random.Next(Alphabet.Length)];
		}
		return new string(chars);
	}
}