import os
import socket
import sys
import concurrent.futures
import time
import random


def domain_megabot(domain, no_match_pattern, whois_server):
    retry_count = 0
    while True:
        response = whois_query(whois_server, domain)
        if response:
            if no_match_pattern in response:
                print(f"{domain} AVAILABLE FOR REGISTRATION!")
                with open("domain.txt", "a") as f:
                    f.write(f"{domain}\n")
            else:
                print(f"{domain} NOT AVAILABLE.")
            # Add random delay to avoid triggering rate limits
            time.sleep(random.uniform(1, 5))
            break
        else:
            retry_count += 1
            sleep_time = min(60, 2**retry_count + random.uniform(0, 1))
            print(f"Retrying query for {domain} in {sleep_time:.2f} seconds...")
            time.sleep(sleep_time)


def whois_query(server, query):
    ip = hostname_to_ip(server)
    if not ip:
        print("FAILED TO RESOLVE HOSTNAME")
        return ""

    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(10)
            s.connect((ip, 43))
            s.sendall(f"{query}\r\n".encode())
            response = b""
            while True:
                data = s.recv(4096)
                if not data:
                    break
                response += data
        return response.decode()
    except socket.error as err:
        print(f"Error during WHOIS query: {err}")
        return None


def hostname_to_ip(hostname):
    try:
        return socket.gethostbyname(hostname)
    except socket.error as err:
        print(f"Error resolving hostname: {err}")
        return None


def str_split(s, delimiter):
    return s.split(delimiter)


def str_conn(s1, s2):
    return f"{s1}{s2}"


def main():
    # Automatically read TLDs from the TLD_DATA file in the current directory
    try:
        with open("TLD_DATA", "r") as fp:
            tld_data = fp.readlines()
    except FileNotFoundError:
        print("TLD DATABASE NOT FOUND!")
        sys.exit(1)

    # Automatically read domain extensions (TLDs) from the DOMAINS file
    try:
        with open("DOMAINS", "r") as fp:
            exts = [line.strip() for line in fp.readlines() if line.strip()]
    except FileNotFoundError:
        print("DOMAINS FILE NOT FOUND!")
        sys.exit(1)

    # Load TLD information
    tld_info = {}
    for line in tld_data:
        arr = str_split(line.strip(), "=")
        if arr[0] in exts:
            tld_info[arr[0]] = {
                "whois_server": arr[1],
                "no_match_pattern": arr[2],
            }

    if not tld_info:
        print("NO VALID TLDs FOUND IN DOMAINS FILE!")
        sys.exit(2)

    # Automatically read the dictionary from the DICT file
    try:
        with open("DICT", "r") as fp_dict:
            dictionary = fp_dict.readlines()
    except FileNotFoundError:
        print("DICTIONARY FILE NOT FOUND!")
        sys.exit(3)

    with open("domain.txt", "w") as fp_result:
        fp_result.write(
            "AVAILABLE DOMAIN NAME LIST:\n" "---------------------------------\n"
        )

    # Prepare domain names to check
    domains = []
    for line in dictionary:
        domain_prefix = line.strip()
        if domain_prefix:
            for ext in exts:
                domain = str_conn(str_conn(domain_prefix, "."), ext)
                domains.append(
                    (
                        domain,
                        tld_info[ext]["no_match_pattern"],
                        tld_info[ext]["whois_server"],
                    )
                )

    # Use a thread pool to check multiple domains concurrently
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = [
            executor.submit(domain_megabot, domain, pattern, server)
            for domain, pattern, server in domains
        ]
        for future in concurrent.futures.as_completed(futures):
            future.result()  # Handle any exceptions in threads

    print("TASK FINISHED! RESULTS SAVED TO domain.txt")


if __name__ == "__main__":
    main()