using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using QRCoder;

/// <summary>
/// Disclaimer: Most of the mDNS related codebase is generated by ChatGPT.
/// </summary>
public class AndroidWirelessDebuggingAutoConnect
{
	private const string ADB_PATH = "adb.exe";

	/// <summary>
	/// If <c>true</c> and an Android device broadcasts a connection service, app will attempt to connect to that device without waiting for
	/// the pair service (which is broadcasted after the QR code is read). This is useful for quickly connecting to the last paired device
	/// without requiring the QR code to be read.
	/// </summary>
	private static bool canConnectWithoutPair = true;
	private static bool hasPaired;

	private class ServiceData
	{
		public readonly string Name;
		public readonly ushort Port;

		public ServiceData( string name, ushort port )
		{
			Name = name;
			Port = port;
		}
	}

	public static void Main( string[] args )
	{
		string name = GetRandomText( 5 );
		string password = GetRandomText( 6 );
		string pairServiceName = $"{name}_adb-tls-pairing";
		string connectServiceName = $"_adb-tls-connect._tcp.local";
		string qrCode = $"WIFI:T:ADB;S:{name};P:{password};;";

		Console.WriteLine( AsciiQRCodeHelper.GetQRCode( qrCode, QRCodeGenerator.ECCLevel.M, invert: false ) );

		using UdpClient client = new UdpClient();
		IPEndPoint localEp = new IPEndPoint( IPAddress.Any, 5353 );

		client.Client.SetSocketOption( SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true );
		client.ExclusiveAddressUse = false;
		client.Client.Bind( localEp );

		IPAddress multicastAddress = IPAddress.Parse( "224.0.0.251" );
		client.JoinMulticastGroup( multicastAddress );

		Console.WriteLine( "Scan the QR code with Wireless debugging..." );

		while( true )
		{
			try
			{
				if( hasPaired || canConnectWithoutPair )
					SendMdnsQuery( client, connectServiceName );

				byte[] data = client.Receive( ref localEp );
				List<ServiceData> discoveredServices = ParseMdnsPacket( data );

				if( !hasPaired && discoveredServices.Find( ( e ) => e.Name.Contains( pairServiceName ) ) is ServiceData pairService )
				{
					Console.WriteLine( $"Pairing with {localEp.Address}:{pairService.Port}..." );

					if( RunADBCommand( $"pair {localEp.Address}:{pairService.Port} {password}" ) )
					{
						hasPaired = true;
						client.Client.ReceiveTimeout = 5000;
						Console.WriteLine( "Fetching connection information..." );
					}
				}
				else if( ( hasPaired || canConnectWithoutPair ) && discoveredServices.Find( ( e ) => e.Name.Contains( connectServiceName ) ) is ServiceData connectService )
				{
					Console.WriteLine( $"Connecting to {localEp.Address}:{connectService.Port}..." );

					if( RunADBCommand( $"connect {localEp.Address}:{connectService.Port}" ) )
					{
						Console.WriteLine( "Connected to device! Press any key to exit..." );
						Console.ReadKey();

						return;
					}
					else if( canConnectWithoutPair )
					{
						Console.WriteLine( "Auto-connect failed. Scan the QR code to connect manually..." );
						canConnectWithoutPair = false;
					}
				}

				Thread.Sleep( 500 );
			}
			catch( TimeoutException )
			{
			}
		}
	}

	private static bool RunADBCommand( string command )
	{
		try
		{
			using( Process adb = new Process() )
			{
				adb.StartInfo = new ProcessStartInfo()
				{
					FileName = ADB_PATH,
					Arguments = command,
					UseShellExecute = false,
					CreateNoWindow = true,
					RedirectStandardOutput = true,
				};

				adb.Start();

				string output = adb.StandardOutput.ReadToEnd();
				Console.WriteLine( output );

				adb.WaitForExit();
				return adb.ExitCode == 0 && output.IndexOf( "failed", StringComparison.OrdinalIgnoreCase ) < 0;
			}
		}
		catch( Exception e )
		{
			Console.WriteLine( $"Error while running adb command: {e}" );
			return false;
		}
	}

	private static string GetRandomText( int length )
	{
		const string randomCharacters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";

		StringBuilder sb = new StringBuilder( length );
		for( int i = 0; i < length; i++ )
			sb.Append( randomCharacters[Random.Shared.Next( randomCharacters.Length )] );

		return sb.ToString();
	}

	private static void SendMdnsQuery( UdpClient udpClient, string serviceName )
	{
		byte[] query = ConstructMdnsQuery( serviceName );
		IPEndPoint endPoint = new IPEndPoint( IPAddress.Parse( "224.0.0.251" ), 5353 );
		udpClient.Send( query, query.Length, endPoint );
	}

	private static byte[] ConstructMdnsQuery( string serviceName )
	{
		// Header section
		byte[] header = new byte[12];
		header[5] = 1; // QDCOUNT (1 question)

		// Question section
		byte[] nameBytes = EncodeDomainName( serviceName );
		byte[] question = new byte[nameBytes.Length + 4];
		Buffer.BlockCopy( nameBytes, 0, question, 0, nameBytes.Length );
		question[nameBytes.Length] = 0; // QTYPE (A record)
		question[nameBytes.Length + 1] = 255; // PTR
		question[nameBytes.Length + 2] = 0; // QCLASS (IN)
		question[nameBytes.Length + 3] = 255;

		// Combine header and question sections
		byte[] query = new byte[header.Length + question.Length];
		Buffer.BlockCopy( header, 0, query, 0, header.Length );
		Buffer.BlockCopy( question, 0, query, header.Length, question.Length );

		return query;
	}

	private static byte[] EncodeDomainName( string domainName )
	{
		var parts = domainName.Split( '.' );
		var result = new byte[domainName.Length + 2]; // Include the length bytes and the final zero byte
		int index = 0;

		foreach( var part in parts )
		{
			result[index++] = (byte) part.Length;
			byte[] partBytes = Encoding.UTF8.GetBytes( part );
			Buffer.BlockCopy( partBytes, 0, result, index, partBytes.Length );
			index += partBytes.Length;
		}
		result[index] = 0; // Null terminator

		return result;
	}

	private static List<ServiceData> ParseMdnsPacket( byte[] data )
	{
		List<ServiceData> result = new List<ServiceData>( 1 );
		try
		{
			int index = 0;

			// Parse DNS header
			ushort transactionId = ReadUInt16( data, ref index );
			ushort flags = ReadUInt16( data, ref index );
			ushort questionCount = ReadUInt16( data, ref index );
			ushort answerCount = ReadUInt16( data, ref index );
			ushort authorityCount = ReadUInt16( data, ref index );
			ushort additionalCount = ReadUInt16( data, ref index );
			//Console.WriteLine( $"Transaction ID: {transactionId} Flags: {flags} Questions: {questionCount} Answers: {answerCount} Authority RRs: {authorityCount} Additional RRs: {additionalCount}" );

			// Parse DNS questions
			for( int i = 0; i < questionCount; i++ )
			{
				string questionName = ReadDomainName( data, ref index );
				ushort questionType = ReadUInt16( data, ref index );
				ushort questionClass = ReadUInt16( data, ref index );
				//Console.WriteLine( $"Question {i + 1} Name: {questionName} Type: {questionType} Class: {questionClass}" );
			}

			// Parse answers
			for( int i = 0; i < answerCount; i++ )
			{
				string answerName = ReadDomainName( data, ref index );
				ushort answerType = ReadUInt16( data, ref index );
				ushort answerClass = ReadUInt16( data, ref index );
				uint ttl = ReadUInt32( data, ref index );
				ushort dataLength = ReadUInt16( data, ref index );

				if( answerType == 33 ) // SRV record
				{
					ushort priority = ReadUInt16( data, ref index );
					ushort weight = ReadUInt16( data, ref index );
					ushort port = ReadUInt16( data, ref index );
					string target = ReadDomainName( data, ref index );
					//Console.WriteLine( $"Answer SRV: {answerName}, Port: {port}, Target: {target}, TTL: {ttl}" );

					result.Add( new ServiceData( answerName, port ) );
				}
				else
				{
					// Skip the rest of the data
					index += dataLength;
				}
			}

			for( int i = 0; i < authorityCount; i++ )
			{
				string name = ReadDomainName( data, ref index );
				ushort type = ReadUInt16( data, ref index );
				ushort classValue = ReadUInt16( data, ref index );
				uint ttl = ReadUInt32( data, ref index );
				ushort dataLength = ReadUInt16( data, ref index );

				if( type == 2 ) // NS record
				{
					string nsDomain = ReadDomainName( data, ref index );
					//Console.WriteLine( $"Authority NS: {name}, NS: {nsDomain}, TTL: {ttl}" );
				}
				else if( type == 6 ) // SOA record
				{
					string mName = ReadDomainName( data, ref index );
					string rName = ReadDomainName( data, ref index );
					uint serial = ReadUInt32( data, ref index );
					uint refresh = ReadUInt32( data, ref index );
					uint retry = ReadUInt32( data, ref index );
					uint expire = ReadUInt32( data, ref index );
					uint minimum = ReadUInt32( data, ref index );
					//Console.WriteLine( $"Authority SOA: {name}, MNAME: {mName}, RNAME: {rName}, Serial: {serial}, Refresh: {refresh}, Retry: {retry}, Expire: {expire}, Minimum: {minimum}" );
				}
				else if( type == 33 ) // SRV record
				{
					ushort priority = ReadUInt16( data, ref index );
					ushort weight = ReadUInt16( data, ref index );
					ushort port = ReadUInt16( data, ref index );
					string target = ReadDomainName( data, ref index );
					//Console.WriteLine( $"Authority SRV: {name}, Port: {port}, Target: {target}, TTL: {ttl}" );

					result.Add( new ServiceData( name, port ) );
				}
				else
				{
					// Skip the rest of the data
					index += dataLength;
				}
			}

			for( int i = 0; i < additionalCount; i++ )
			{
				string name = ReadDomainName( data, ref index );
				ushort type = ReadUInt16( data, ref index );
				ushort classValue = ReadUInt16( data, ref index );
				uint ttl = ReadUInt32( data, ref index );
				ushort dataLength = ReadUInt16( data, ref index );

				if( type == 33 ) // SRV record
				{
					ushort priority = ReadUInt16( data, ref index );
					ushort weight = ReadUInt16( data, ref index );
					ushort port = ReadUInt16( data, ref index );
					string target = ReadDomainName( data, ref index );
					//Console.WriteLine( $"Additional SRV: {name}, Port: {port}, Target: {target}, TTL: {ttl}" );

					result.Add( new ServiceData( name, port ) );
				}
				else
				{
					// Skip the rest of the data
					index += dataLength;
				}
			}
		}
		catch( Exception ex )
		{
			Console.WriteLine( $"Error parsing mDNS packet: {ex}" );
		}

		return result;
	}

	private static string ReadDomainName( byte[] data, ref int index )
	{
		StringBuilder name = new StringBuilder();
		while( data[index] != 0 )
		{
			byte length = data[index++];
			if( ( length & 0xC0 ) == 0xC0 ) // Compression
			{
				ushort pointer = (ushort) ( ( ( length & 0x3F ) << 8 ) | data[index++] );
				int savedIndex = index;
				index = pointer;
				name.Append( ReadDomainName( data, ref index ) );
				index = savedIndex;
				return name.ToString();
			}

			if( name.Length > 0 )
				name.Append( '.' );

			name.Append( Encoding.UTF8.GetString( data, index, length ) );
			index += length;
		}

		index++; // Skip the null byte at the end of the domain name
		return name.ToString();
	}

	private static ushort ReadUInt16( byte[] data, ref int index )
	{
		ushort value = (ushort) ( ( data[index] << 8 ) | data[index + 1] );
		index += 2;
		return value;
	}

	private static uint ReadUInt32( byte[] data, ref int index )
	{
		uint value = (uint) ( ( data[index] << 24 ) | ( data[index + 1] << 16 ) | ( data[index + 2] << 8 ) | data[index + 3] );
		index += 4;
		return value;
	}
}