#!/usr/bin/env python3
"""Smoke test for testserver.py — every hostile mode answers the way its contract says.

Not a substitute for CORE's own conformance against it; this just guards the server from
regressing. Run directly (`tools/testserver/selftest.py`) or via ctest.
"""

from __future__ import annotations

import hashlib
import http.client
import json
import socket
import subprocess
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
SERVER = HERE / "testserver.py"
SEED = 42

_fail = 0


def check(name: str, cond: bool, detail: str = "") -> None:
    global _fail
    mark = "\033[32mok\033[0m" if cond else "\033[31mFAIL\033[0m"
    print(f"  {mark}  {name}" + (f"  — {detail}" if detail and not cond else ""))
    if not cond:
        _fail += 1


class Server:
    def __init__(self) -> None:
        self.proc = subprocess.Popen(
            [sys.executable, str(SERVER), "--port", "0", "--seed", str(SEED),
             "--loris-seconds", "1", "--throttle-bps", str(256 * 1024)],
            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True,
        )
        self.port = int(self.proc.stdout.readline().strip())
        for _ in range(50):
            try:
                with socket.create_connection(("127.0.0.1", self.port), timeout=0.2):
                    break
            except OSError:
                time.sleep(0.05)

    def conn(self) -> http.client.HTTPConnection:
        return http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)

    def get(self, path: str, headers: dict | None = None):
        c = self.conn()
        c.request("GET", path, headers=headers or {})
        r = c.getresponse()
        body = r.read()
        c.close()
        return r, body

    def raw_get(self, path: str, headers: dict | None = None) -> tuple[int, dict, bytes, bool]:
        """Returns (status, headers, body_bytes, clean_eof). Tolerates a mid-body RST."""
        s = socket.create_connection(("127.0.0.1", self.port), timeout=10)
        req = f"GET {path} HTTP/1.1\r\nHost: 127.0.0.1:{self.port}\r\nConnection: close\r\n"
        for k, v in (headers or {}).items():
            req += f"{k}: {v}\r\n"
        req += "\r\n"
        s.sendall(req.encode())
        chunks = []
        clean = True
        try:
            while True:
                b = s.recv(65536)
                if not b:
                    break
                chunks.append(b)
        except ConnectionResetError:
            clean = False
        finally:
            s.close()
        raw = b"".join(chunks)
        head, _, body = raw.partition(b"\r\n\r\n")
        lines = head.split(b"\r\n")
        status = int(lines[0].split()[1]) if lines and lines[0] else 0
        hdrs = {}
        for ln in lines[1:]:
            if b":" in ln:
                k, v = ln.split(b":", 1)
                hdrs[k.decode().strip().lower()] = v.decode().strip()
        return status, hdrs, body, clean

    def close(self) -> None:
        self.proc.terminate()
        try:
            self.proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            self.proc.kill()


def ref_sha(srv: Server, path_mode: str, size_spec: str) -> tuple[str, int]:
    r, body = srv.get(f"/{path_mode}/sha256/{size_spec}")
    j = json.loads(body)
    return j["sha256"], j["size"]


def main() -> int:
    srv = Server()
    print(f"testserver on :{srv.port}  seed={SEED}")
    try:
        # health
        r, body = srv.get("/__health")
        check("health", r.status == 200 and body == b"ok")

        # plain: full body matches the reference hash, Range works, 206 + Content-Range
        want, size = ref_sha(srv, "plain", "256K")
        r, body = srv.get("/plain/file/256K")
        check("plain full 200", r.status == 200 and len(body) == size)
        check("plain sha256 matches", hashlib.sha256(body).hexdigest() == want)
        check("plain advertises Accept-Ranges", r.getheader("Accept-Ranges") == "bytes")
        r, body = srv.get("/plain/file/256K", {"Range": "bytes=1000-1999"})
        check("plain range 206", r.status == 206 and len(body) == 1000)
        check("plain Content-Range", r.getheader("Content-Range") == f"bytes 1000-1999/{size}")
        check("plain range bytes correct", body == full_slice(srv, "plain", "256K", 1000, 2000))

        # no-range: no Accept-Ranges, Range ignored -> 200 full
        r, body = srv.get("/no-range/file/128K", {"Range": "bytes=0-1023"})
        check("no-range ignores Range (200)", r.status == 200 and len(body) == 128 * 1024)
        check("no-range hides Accept-Ranges", r.getheader("Accept-Ranges") is None)

        # lies-about-accept-ranges: advertises, ignores
        r, body = srv.get("/lies-about-accept-ranges/file/128K", {"Range": "bytes=0-1023"})
        check("lies advertises Accept-Ranges", r.getheader("Accept-Ranges") == "bytes")
        check("lies still sends 200 full", r.status == 200 and len(body) == 128 * 1024)

        # etag-changes: ETag differs across requests; If-Range mismatch -> 200
        r1, _ = srv.get("/etag-changes/file/64K")
        r2, _ = srv.get("/etag-changes/file/64K")
        check("etag-changes: ETag varies", r1.getheader("ETag") != r2.getheader("ETag"))
        r, body = srv.get("/etag-changes/file/64K",
                          {"Range": "bytes=0-99", "If-Range": '"stale"'})
        check("etag-changes: If-Range mismatch -> 200 full", r.status == 200 and len(body) == 64 * 1024)

        # 416-always: Range -> 416 with Content-Range */size; plain GET -> 200
        r, body = srv.get("/416-always/file/64K", {"Range": "bytes=0-99"})
        check("416-always range -> 416", r.status == 416)
        check("416-always Content-Range */n", r.getheader("Content-Range") == f"bytes */{64*1024}")
        r, body = srv.get("/416-always/file/64K")
        check("416-always plain GET -> 200", r.status == 200 and len(body) == 64 * 1024)

        # redirect-chain: 302s then 200 (http.client doesn't auto-follow; do it by hand)
        status, body, hops = follow(srv, "/redirect-chain/file/16K")
        check("redirect-chain lands 200", status == 200 and len(body) == 16 * 1024)
        check("redirect-chain actually bounced", hops >= 3, f"{hops} hops")

        # 401-basic
        r, _ = srv.get("/401-basic/file/16K")
        check("401-basic challenges", r.status == 401 and "Basic" in (r.getheader("WWW-Authenticate") or ""))
        import base64 as _b64
        tok = _b64.b64encode(b"test:test").decode()
        r, body = srv.get("/401-basic/file/16K", {"Authorization": f"Basic {tok}"})
        check("401-basic accepts test:test", r.status == 200 and len(body) == 16 * 1024)

        # 403-without-referer
        r, _ = srv.get("/403-without-referer/file/16K")
        check("403 without Referer", r.status == 403)
        r, body = srv.get("/403-without-referer/file/16K",
                          {"Referer": f"http://127.0.0.1:{srv.port}/page"})
        check("403 mode ok with Referer", r.status == 200 and len(body) == 16 * 1024)

        # content-length-mismatch: header length > bytes delivered, unclean close
        st, hdrs, body, clean = srv.raw_get("/content-length-mismatch/file/64K")
        check("clen-mismatch: declared > received",
              int(hdrs.get("content-length", "0")) > len(body), f"{hdrs.get('content-length')} vs {len(body)}")
        check("clen-mismatch: connection not clean", clean is False)

        # flaky-reset: first two attempts cut + RST, third is clean & correct
        want, size = ref_sha(srv, "flaky-reset", "64K")
        results = [srv.raw_get("/flaky-reset/file/64K") for _ in range(3)]
        check("flaky-reset: attempt 1 unclean", results[0][3] is False and len(results[0][2]) < size)
        check("flaky-reset: attempt 2 unclean", results[1][3] is False)
        check("flaky-reset: attempt 3 clean & full", results[2][3] is True and len(results[2][2]) == size)
        check("flaky-reset: attempt 3 hash ok", hashlib.sha256(results[2][2]).hexdigest() == want)

        # chunked-no-length
        st, hdrs, body, clean = srv.raw_get("/chunked-no-length/file/32K")
        check("chunked: TE chunked, no CL",
              hdrs.get("transfer-encoding") == "chunked" and "content-length" not in hdrs)
        dechunked = dechunk(body)
        want, size = ref_sha(srv, "chunked-no-length", "32K")
        check("chunked: body decodes to full file", len(dechunked) == size)
        check("chunked: hash ok", hashlib.sha256(dechunked).hexdigest() == want)

        # throttled: 256K at 256 KiB/s budget -> takes >= ~0.9s
        t0 = time.monotonic()
        r, body = srv.get("/throttled/file/512K")
        dt = time.monotonic() - t0
        check("throttled: rate-limited", dt >= 0.9 and len(body) == 512 * 1024, f"{dt:.2f}s")

        # slow-loris: eventually completes, slowly
        t0 = time.monotonic()
        st, hdrs, body, clean = srv.raw_get("/slow-loris/file/8K")
        dt = time.monotonic() - t0
        check("slow-loris: completes", len(body) == 8 * 1024, f"got {len(body)}")
        check("slow-loris: was slow", dt >= 0.5, f"{dt:.2f}s")

        # utf8 / legacy content-disposition
        r, _ = srv.get("/utf8-content-disposition/file/8K")
        cd = r.getheader("Content-Disposition") or ""
        check("utf8 CD has filename*=UTF-8''", "filename*=UTF-8''%E2%82%AC" in cd)
        r, _ = srv.get("/legacy-content-disposition/file/8K")
        cd = r.getheader("Content-Disposition") or ""
        check("legacy CD has encoded-word", "=?UTF-8?B?" in cd)

        # expiring-signed-url
        r, body = srv.get("/expiring-signed-url/sign/8K?ttl=2")
        signed = json.loads(body)["url"]
        path = signed.split(f":{srv.port}", 1)[1]
        r, body = srv.get(path)
        check("signed URL works before expiry", r.status == 200 and len(body) == 8 * 1024)
        r, _ = srv.get("/expiring-signed-url/file/8K")
        check("unsigned request rejected", r.status == 403)
        time.sleep(2.1)
        r, body = srv.get(path)
        check("signed URL rejected after expiry", r.status == 403 and b"expired" in body)

        # __control reset: prime the counter to attempt 2, reset, next attempt is 1 again
        # (unclean) rather than 3 (which would be clean).
        srv.raw_get("/flaky-reset/file/4K")   # attempt 1
        srv.raw_get("/flaky-reset/file/4K")   # attempt 2
        _, _, _, clean_before = srv.raw_get("/flaky-reset/file/4K")  # attempt 3 -> clean
        check("flaky-reset: 3rd attempt would be clean", clean_before is True)
        srv.raw_get("/flaky-reset/file/4K")   # attempt 1 again after this reset...
        c = srv.conn()
        payload = b'{"reset":true}'
        c.request("POST", "/__control", body=payload,
                  headers={"Content-Length": str(len(payload))})
        rr = c.getresponse()
        rr.read()
        c.close()
        check("__control reset 200", rr.status == 200)
        _, _, _, clean1 = srv.raw_get("/flaky-reset/file/4K")  # attempt 1 post-reset
        _, _, _, clean2 = srv.raw_get("/flaky-reset/file/4K")  # attempt 2 post-reset
        check("__control reset the counter (1st post-reset unclean)", clean1 is False)
        check("__control reset the counter (2nd post-reset unclean)", clean2 is False)

    finally:
        srv.close()

    print()
    if _fail:
        print(f"\033[31m{_fail} check(s) failed\033[0m")
        return 1
    print("\033[32mall checks passed\033[0m")
    return 0


def full_slice(srv: Server, mode: str, size_spec: str, lo: int, hi: int) -> bytes:
    r, body = srv.get(f"/{mode}/file/{size_spec}", {"Range": f"bytes={lo}-{hi-1}"})
    return body


def follow(srv: Server, path: str, limit: int = 10) -> tuple[int, bytes, int]:
    """Manually follow 3xx Location headers. Returns (final_status, body, hop_count)."""
    hops = 0
    for _ in range(limit):
        r, body = srv.get(path)
        if r.status in (301, 302, 303, 307, 308):
            loc = r.getheader("Location") or ""
            path = loc.split(f":{srv.port}", 1)[1] if f":{srv.port}" in loc else loc
            hops += 1
            continue
        return r.status, body, hops
    return 0, b"", hops


def dechunk(raw: bytes) -> bytes:
    out = bytearray()
    i = 0
    while i < len(raw):
        j = raw.find(b"\r\n", i)
        if j < 0:
            break
        n = int(raw[i:j].split(b";")[0], 16)
        if n == 0:
            break
        out += raw[j + 2:j + 2 + n]
        i = j + 2 + n + 2
    return bytes(out)


if __name__ == "__main__":
    raise SystemExit(main())