using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Runtime.InteropServices;
using System.Text;

namespace DnsForwarder
{
    internal class Program
    {
        private static string DefaultServer;
        private static string ChinaServer;
        private static Dictionary<string, string> ChinaList;

        private static void Main(string[] args)
        {
            if (args.Length != 3)
            {
                Console.WriteLine("<default dns server ip> <china dns server ip> <path to dnsmasq-china-list file>");
                Environment.Exit(-1);
                return;
            }

            Console.WriteLine("Starting...");
            Console.WriteLine($"Default Server:{DefaultServer = args[0]}");
            Console.WriteLine($"China DNS Server:{ChinaServer = args[1]}");
            Console.WriteLine($"dnsmasq-china-list file:{args[2]}");

            Console.WriteLine("Loading list file...");
            ChinaList = LoadListFile(args[2]).ToDictionary(p => p[0], p => p[1]);

            Console.WriteLine("OK. Listing...");

            var listen = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            listen.Bind(new IPEndPoint(IPAddress.Any, 53));

            while (true)
            {
                EndPoint inEp = new IPEndPoint(IPAddress.Any, 0);
                var buf = new byte[1500];
                var count = listen.ReceiveFrom(buf, 0, buf.Length, SocketFlags.None, ref inEp);
                var domain = ExtractDomainName(buf, count);

                var target = MatchServer(domain);

                Console.WriteLine($"{DateTime.Now:yyyyMMdd HH:mm:ss} {inEp} [{target}]\t{domain}");

                listen.SendTo(GetDnsResponse(buf, count, target), inEp);
            }
        }

        private static string ExtractDomainName(byte[] buf, int count)
        {
            var lst = new List<byte[]>();

            var ptr = 12;
            while (buf[ptr] != 0)
            {
                ptr += buf[ptr] + 1;
            }

            var bufDomain = new byte[ptr - 12];
            Array.Copy(buf, 12, bufDomain, 0, bufDomain.Length);

            //dot fill
            ptr = 0;
            while (ptr < bufDomain.Length)
            {
                var b = bufDomain[ptr];
                bufDomain[ptr] = (byte)'.';
                ptr += b + 1;
            }

            return Encoding.ASCII.GetString(bufDomain, 1, bufDomain.Length - 1);
        }

        private static byte[] GetDnsResponse(byte[] buf, int count, string host, int port = 53)
        {
            using var to = new UdpClient();
            to.Connect(host, 53);
            to.Send(buf, count);
            IPEndPoint outEp = null;
            return to.Receive(ref outEp);
        }

        private static string MatchServer(string domain)
        {
            var lower = domain.ToLower();
            if (lower.EndsWith(".cn")) return ChinaServer;

            var parts = lower.Split('.').Reverse().ToArray();

            for (int i = parts.Length - 1; i >= 0; i--)
            {
                var d = string.Join(".", parts.Take(i + 1).Reverse());
                if (ChinaList.TryGetValue(d, out var tar)) return tar;
            }
            
            return DefaultServer;
        }

        private static string[][] LoadListFile(string path)
        {
            var lines = File.ReadAllLines(path);
            var items = lines.Select(p =>
            {
                if (p.StartsWith("#")) return null;
                var parts = p.Split('/');
                return new[] { parts[1], parts[2] };
            }).Where(p => p != null).ToArray();
            return items;
        }
    }
}