import gradio as gr
import requests, time, socket, ssl, random
from ping3 import ping
import dns.resolver
from bs4 import BeautifulSoup


def dns_lookup(domain):
    try:
        start = time.time()
        dns.resolver.resolve(domain)
        return round((time.time() - start) * 1000, 2)
    except:
        return None


def ssl_handshake(domain):
    try:
        ctx = ssl.create_default_context()
        sock = socket.create_connection((domain, 443), timeout=10)
        start = time.time()
        ssock = ctx.wrap_socket(sock, server_hostname=domain)
        t = (time.time() - start) * 1000
        ssock.close()
        return round(t, 2)
    except:
        return None


def test_speed(url):
    if not url.startswith("http"):
        url = "http://" + url

    domain = url.split("//")[1].split("/")[0]
    result = {}

    # DNS
    dns_time = dns_lookup(domain)
    result["DNS Lookup (ms)"] = dns_time

    # Ping
    p = ping(domain, timeout=1)
    result["Ping (ms)"] = round(p * 1000, 2) if p else None

    # SSL
    ssl_time = ssl_handshake(domain)
    result["SSL Handshake (ms)"] = ssl_time

    try:
        session = requests.Session()
        cache_bypass_url = url + "?" + str(random.randint(100000, 999999))

        start_ttfb = time.time()
        r = session.get(cache_bypass_url, stream=True, timeout=20)
        ttfb = (time.time() - start_ttfb) * 1000

        result["HTTP Status"] = r.status_code
        result["TTFB (ms)"] = round(ttfb, 2)
        result["Final URL"] = r.url
        result["Redirect Chain"] = [resp.url for resp in r.history] + [r.url]

        # Download
        start_dl = time.time()
        content = r.content
        download_time = (time.time() - start_dl) * 1000

        result["Download Time (ms)"] = round(download_time, 2)
        result["Page Size (KB)"] = round(len(content) / 1024, 2)

        # ---- ASSET COUNT ----
        soup = BeautifulSoup(content, "html.parser")

        images = soup.find_all("img")
        videos = soup.find_all("video")
        audios = soup.find_all("audio")

        result["Images Count"] = len(images)
        result["Videos Count"] = len(videos)
        result["Audios Count"] = len(audios)
        result["Total Assets"] = len(images) + len(videos) + len(audios)

        # Speed index
        speed_index = sum(
            x for x in [dns_time, ssl_time, ttfb, download_time * 0.5] if x
        )
        result["Speed Index (ms)"] = round(speed_index, 2)
        result["Speed Index (s)"] = round(speed_index / 1000, 2)

    except Exception as e:
        return {"Error": str(e)}

    # Rating
    if ttfb < 200:
        rating = "Fast"
    elif ttfb < 600:
        rating = "Medium"
    else:
        rating = "Slow"

    result["Speed Rating"] = rating
    return result


ui = gr.Interface(
    fn=test_speed,
    inputs=gr.Textbox(label="Enter URL"),
    outputs=gr.JSON(label="Result"),
    title="Complete Website Speed Tester",
)

if __name__ == "__main__":
    ui.launch()