#!/usr/bin/env python3
"""phases.py — measure per-phase HTTP latency using raw sockets.

Phases: DNS lookup, TCP connect, TLS handshake (HTTPS only),
Server/TTFB (sent → first byte), Transfer (first byte → EOF), Total.

Usage:
    python phases.py <url> [url ...]          # one or more URLs
    python phases.py <config_file>            # plain-text list of URLs

Config file: one URL per line; '#' lines and blank lines ignored.
Exits 0 if all URLs complete without network error, 1 if any failed.
"""

import socket
import ssl
import sys
import time
import urllib.parse
from dataclasses import dataclass, field

_TIMEOUT = 10.0
_SEP = "─" * 29

# (dataclass field name, 14-char display label)
_PHASE_LABELS = [
    ("dns",      "DNS lookup    "),
    ("connect",  "TCP connect   "),
    ("tls",      "TLS handshake "),
    ("ttfb",     "Server (TTFB) "),
    ("transfer", "Transfer      "),
]


@dataclass
class Phase:
    ms: float = 0.0
    present: bool = False


@dataclass
class Result:
    url: str
    dns:      Phase = field(default_factory=Phase)
    connect:  Phase = field(default_factory=Phase)
    tls:      Phase = field(default_factory=Phase)
    ttfb:     Phase = field(default_factory=Phase)
    transfer: Phase = field(default_factory=Phase)
    total:    Phase = field(default_factory=Phase)
    status_code: int = 0
    fail_phase: str = ""
    err: Exception | None = None


def _phase(start: float, end: float) -> Phase:
    return Phase(ms=(end - start) * 1000, present=True)


def _parse_status(data: bytes) -> int:
    """Extract HTTP status code from the first recv() chunk."""
    eol = data.find(b"\r\n")
    if eol == -1:
        eol = data.find(b"\n")
    if eol == -1:
        return 0
    parts = data[:eol].decode("latin-1", errors="replace").split(None, 2)
    if len(parts) >= 2:
        try:
            return int(parts[1])
        except ValueError:
            pass
    return 0


def measure(raw_url: str) -> Result:
    """Probe raw_url and return a Result with per-phase timings.

    Partial phases are preserved when the request fails mid-flight.
    Method is always GET; bodies are fully drained so Transfer timing is real.
    Redirects are NOT followed — the raw HTTP response status is reported.
    """
    r = Result(url=raw_url)

    parsed = urllib.parse.urlparse(raw_url)
    scheme = parsed.scheme.lower()
    if scheme not in ("http", "https"):
        r.fail_phase = "request"
        r.err = ValueError(f"unsupported scheme: {scheme!r}")
        return r

    host = parsed.hostname or ""
    port = parsed.port or (443 if scheme == "https" else 80)
    path = parsed.path or "/"
    if parsed.query:
        path = path + "?" + parsed.query
    use_tls = scheme == "https"

    t_start = time.perf_counter()

    # ── DNS ──────────────────────────────────────────────────────────────────

    t0 = time.perf_counter()
    try:
        infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
    except socket.gaierror as exc:
        r.fail_phase = "dns"
        r.err = exc
        r.total = _phase(t_start, time.perf_counter())
        return r
    r.dns = _phase(t0, time.perf_counter())

    # ── TCP connect ───────────────────────────────────────────────────────────

    addr = infos[0][4]
    family = infos[0][0]
    sock = socket.socket(family, socket.SOCK_STREAM)
    sock.settimeout(_TIMEOUT)
    t0 = time.perf_counter()
    try:
        sock.connect(addr)
    except socket.timeout as exc:
        sock.close()
        r.connect = _phase(t0, time.perf_counter())
        r.fail_phase = "timeout"
        r.err = exc
        r.total = _phase(t_start, time.perf_counter())
        return r
    except OSError as exc:
        sock.close()
        r.connect = _phase(t0, time.perf_counter())
        r.fail_phase = "connect"
        r.err = exc
        r.total = _phase(t_start, time.perf_counter())
        return r
    r.connect = _phase(t0, time.perf_counter())

    # ── TLS handshake (HTTPS only) ────────────────────────────────────────────

    if use_tls:
        ctx = ssl.create_default_context()
        t0 = time.perf_counter()
        try:
            # wrap_socket blocks until the handshake completes (do_handshake_on_connect=True)
            sock = ctx.wrap_socket(sock, server_hostname=host)
        except socket.timeout as exc:
            # Timeout during handshake classifies as timeout, not tls (mirrors Go)
            r.tls = _phase(t0, time.perf_counter())
            r.fail_phase = "timeout"
            r.err = exc
            r.total = _phase(t_start, time.perf_counter())
            return r
        except (ssl.SSLError, OSError) as exc:
            r.tls = _phase(t0, time.perf_counter())
            r.fail_phase = "tls"
            r.err = exc
            r.total = _phase(t_start, time.perf_counter())
            return r
        r.tls = _phase(t0, time.perf_counter())

    # ── Send request ──────────────────────────────────────────────────────────

    request = (
        f"GET {path} HTTP/1.1\r\n"
        f"Host: {host}\r\n"
        f"Connection: close\r\n"
        f"User-Agent: latprobe-phases/1.0\r\n"
        f"\r\n"
    ).encode()
    try:
        sock.sendall(request)
    except OSError as exc:
        sock.close()
        r.fail_phase = "request"
        r.err = exc
        r.total = _phase(t_start, time.perf_counter())
        return r
    t_wrote = time.perf_counter()

    # ── TTFB — sent → first response byte ────────────────────────────────────

    try:
        first_chunk = b""
        while not first_chunk:
            chunk = sock.recv(4096)
            if not chunk:
                raise OSError("server closed connection before sending a response")
            first_chunk = chunk
    except socket.timeout as exc:
        sock.close()
        r.fail_phase = "timeout"
        r.err = exc
        r.total = _phase(t_start, time.perf_counter())
        return r
    except OSError as exc:
        sock.close()
        r.fail_phase = "transfer"
        r.err = exc
        r.total = _phase(t_start, time.perf_counter())
        return r

    t_first_byte = time.perf_counter()
    r.ttfb = _phase(t_wrote, t_first_byte)
    r.status_code = _parse_status(first_chunk)

    # ── Transfer — drain remaining body ───────────────────────────────────────

    try:
        while True:
            chunk = sock.recv(65536)
            if not chunk:
                break
    except OSError:
        pass  # connection reset during body drain is acceptable; timing is captured
    finally:
        sock.close()

    t_end = time.perf_counter()
    r.transfer = _phase(t_first_byte, t_end)
    r.total = _phase(t_start, t_end)
    return r


def print_result(r: Result) -> None:
    if r.err is not None and r.status_code == 0:
        print(f"{r.url}  (FAILED)")
    else:
        print(f"{r.url}  ({r.status_code})")

    for attr, label in _PHASE_LABELS:
        phase: Phase = getattr(r, attr)
        if phase.present:
            print(f"  {label} :  {phase.ms:8.2f} ms")

    print(f"  {_SEP}")
    if r.total.present:
        print(f"  {'Total         '} :  {r.total.ms:8.2f} ms")
    if r.err is not None:
        print(f"  ✗ {r.fail_phase}: {r.err}")


def load_config(path: str) -> list[str]:
    urls = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            urls.append(line.split()[0])
    return urls


def main(argv: list[str]) -> int:
    if not argv:
        print(__doc__, file=sys.stderr)
        return 1

    if argv[0].startswith(("http://", "https://")):
        urls = argv
    else:
        try:
            urls = load_config(argv[0])
        except FileNotFoundError:
            print(f"error: file not found: {argv[0]}", file=sys.stderr)
            return 1
        except OSError as exc:
            print(f"error: cannot read config: {exc}", file=sys.stderr)
            return 1

    if not urls:
        print("error: no URLs found", file=sys.stderr)
        return 1

    any_failed = False
    for i, url in enumerate(urls):
        if i > 0:
            print()
        r = measure(url)
        print_result(r)
        if r.err is not None:
            any_failed = True

    return 1 if any_failed else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))