import t, c
from stdint import *
import string
import stdio
import stdlib
import requests
import requests.https as rhttps
import vthreading as vt
import win32.win32file as wf
import win32.win32base as wb
import win32.win32sync as wsync


# ============================================================
# requests/download.vp - large-file ranged download (Range + serial/multi-threaded)
#
# Usage:
#     import requests.download as rd
#     rc: INT = rd.download("https://host/big.bin", "big.bin", threads=4,
#                           chunk_size=262144, verify=True, ca_file="ca.crt")
#     if rc != 0:
#         print(rd.last_error())
#
# Strategy:
#   threads == 1  serial Range chunks (fetched and written to the file in order)
#   threads  > 1  multi-threaded: the file is split into bands by thread count, each thread uses one
#                 connection to download its band's chunks in order, writing at the corresponding offsets of the same file.
#
# Notes:
#   - probes total size and server Range support via Content-Range;
#   - http:// goes through requests (pure socket), https:// goes through requests.https (openssl);
#   - each chunk request omits Accept-Encoding, the response is raw bytes and is written straight to disk.
# ============================================================

define DL_MAX_THREADS = 8
define DL_DEFAULT_CHUNK = 262144
define DL_REDIRECT_MAX = 5

_dl_err_store: t.CArray[char, 512]


def last_error() -> str:
    """Last download failure reason (empty string if none)."""
    return _dl_err_store


def _dl_set_err(text: str):
    if text is None:
        return
    _dl_err_store[0] = 0
    stdio.snprintf(_dl_err_store, 512, "%s", text)


def _dl_strdup(s: str) -> str:
    """Duplicate string on stdlib heap (caller frees); None on failure."""
    if s is None:
        return None
    n: sizet = string.strlen(s)
    out: str = stdlib.malloc(n + 1)
    if out is None:
        return None
    if n > 0:
        string.memcpy(out, s, n)
    out[n] = 0
    return out


# ============================================================
# Single-slice fetch (redirect follow, <=5 hops; dispatch by http/https scheme)
# ============================================================

def _dl_slice(url: str, pos: sizet, want: INT, verify: bool, ca_file: str,
              dst: str, dst_cap: INT, out_total: sizet*) -> INT:
    """Fetch slice [pos, pos+want); body bytes written to dst (locate body via _dl_body_off).
    Returns -2 on too many redirects; -3 on network/server failure.
    """
    if url is None or dst is None or out_total is None:
        return -1
    cur: str = _dl_strdup(url)
    if cur is None:
        return -1
    hops: INT = 0
    rc: INT = 0
    while True:
        is_https: INT = 1 if (string.strncmp(cur, "https://", 8) == 0) else 0
        st: INT = 0
        tot: sizet = 0
        loc: t.CArray[char, 512]
        if is_https != 0:
            rc = rhttps._DlOnce(cur, pos, want, verify, ca_file, dst, dst_cap,
                                &st, loc, 512, &tot)
        else:
            rc = requests._DlOnce(cur, pos, want, dst, dst_cap,
                                  &st, loc, 512, &tot)
        if rc == -2 and loc[0] != 0:
            hops += 1
            if hops > DL_REDIRECT_MAX:
                _dl_set_err("too many redirects")
                stdlib.free(cur)
                return -2
            new_url: str = stdlib.malloc(2048)
            if new_url is None:
                stdlib.free(cur)
                return -1
            if requests._ResolveRedirectUrl(cur, loc, new_url, 2048) != 0:
                stdlib.free(new_url)
                stdlib.free(cur)
                _dl_set_err("bad redirect location")
                return -3
            stdlib.free(cur)
            cur = new_url
            continue
        break
    stdlib.free(cur)
    if rc < 0:
        return rc
    *out_total = tot
    return rc


def _dl_body_off(dst: str, dst_cap: INT) -> INT:
    """Offset of response body (after headers) in dst; -1 if not found."""
    return requests._FindBodyStart(dst, dst_cap)


# ============================================================
# File write (write at offset; supports concurrent threads writing distinct ranges)
# ============================================================

def _dl_write_at(path: str, src: str, length: INT, offset: sizet) -> INT:
    """Open file (create if absent), seek to offset, write length bytes. 0 ok; -1 fail."""
    if path is None or src is None or length <= 0:
        return -1
    h: wb.HANDLE = wf.CreateFileA(
        path, wf.GENERIC_WRITE,
        wf.FILE_SHARE_READ | wf.FILE_SHARE_WRITE | wf.FILE_SHARE_DELETE,
        None, wf.OPEN_ALWAYS, wf.FILE_ATTRIBUTE_NORMAL, None)
    if h == wb.INVALID_HANDLE_VALUE:
        return -1
    # Use 32-bit SetFilePointer (lDistanceToMove passed by value as LONG; no struct ABI ambiguity).
    # lpDistanceToMoveHigh=NULL limits distance to 32-bit signed (<2GiB), so guard the upper bound first.
    if offset > sizet(0x7FFFFFFF):
        wb.CloseHandle(h)
        return -1
    ok_seek: LONG = wf.SetFilePointer(h, LONG(offset), None, wf.FILE_BEGIN)
    # INVALID_SET_FILE_POINTER(0xFFFFFFFF) truncates to LONG -1; offset is already clamped to
    # 0..0x7FFFFFFF, so 0xFFFFFFFF cannot be a valid offset and no GetLastError check is needed.
    if ok_seek == -1:
        wb.CloseHandle(h)
        return -1
    written: ULONG = 0
    ok_w: BOOL = wf.WriteFile(h, src, uint32(length), &written, None)
    wb.CloseHandle(h)
    if ok_w == 0:
        return -1
    if int(written) != length:
        return -1
    return 0


def _dl_truncate(path: str) -> INT:
    """Create (truncate) an empty file via CREATE_ALWAYS. 0 ok; -1 fail."""
    h: wb.HANDLE = wf.CreateFileA(
        path, wf.GENERIC_WRITE,
        wf.FILE_SHARE_READ | wf.FILE_SHARE_WRITE | wf.FILE_SHARE_DELETE,
        None, wf.CREATE_ALWAYS, wf.FILE_ATTRIBUTE_NORMAL, None)
    if h == wb.INVALID_HANDLE_VALUE:
        return -1
    wb.CloseHandle(h)
    return 0


# ============================================================
# Probe: total size / Range support (first slice bytes=0-0)
# ============================================================

def _dl_probe(url: str, verify: bool, ca_file: str,
              out_total: sizet*) -> INT:
    """Probe server total size and Range support. Returns 0=Range supported (out_total valid);
    1=server returned 200 (no Range, full download, out_total=Content-Length); -3 fail.
    """
    *out_total = 0
    buf: str = stdlib.malloc(8192)
    if buf is None:
        _dl_set_err("out of memory")
        return -3
    tot: sizet = 0
    rc: INT = _dl_slice(url, 0, 0, verify, ca_file, buf, 8192, &tot)
    stdlib.free(buf)
    if rc < 0:
        _dl_set_err("probe failed")
        return -3
    if tot > 0:
        *out_total = tot
        return 0
    _dl_set_err("cannot determine size")
    return -3


# ============================================================
# Download one band ([b_start, b_end)), fetching and writing slice by slice
# ============================================================

def _dl_download_band(path: str, url: str, is_https: INT, verify: bool, ca_file: str,
                      b_start: sizet, b_end: sizet, chunk: INT,
                      buf: str, buf_cap: INT,
                      p_done: sizet*) -> INT:
    """0 ok; -1 network/server; -2 write failure. p_done accumulates bytes written."""
    pos: sizet = b_start
    while pos < b_end:
        want: INT = chunk
        if b_end - pos < sizet(want):
            want = INT(b_end - pos)
        tot: sizet = 0
        rc: INT = _dl_slice(url, pos, want, verify, ca_file, buf, buf_cap, &tot)
        if rc <= 0:
            _dl_set_err("slice fetch failed")
            return -1
        bo: INT = _dl_body_off(buf, buf_cap)
        if bo < 0:
            _dl_set_err("bad slice response")
            return -1
        got: INT = rc
        if bo + got > buf_cap:
            got = buf_cap - bo
        if got <= 0:
            _dl_set_err("slice empty")
            return -1
        if _dl_write_at(path, buf + bo, got, pos) != 0:
            _dl_set_err("write file failed")
            return -2
        *p_done = *p_done + sizet(got)
        pos += sizet(got)
        if got < want:
            # Server ended early (only the last slice is < want on normal completion)
            if pos >= b_end:
                break
            _dl_set_err("short slice")
            return -1
    return 0


# ============================================================
# Multi-thread worker
# ============================================================
class DlJob:
    path: str
    url: str
    is_https: INT
    verify: INT
    ca_file: str
    start: sizet
    end: sizet
    chunk: INT
    buf: str
    buf_cap: INT
    done: sizet
    status: INT


def _dl_worker(arg: VOIDPTR) -> t.CPtr:
    if arg is None:
        return None
    job: DlJob* = arg
    rc: INT = _dl_download_band(
        job.path, job.url, job.is_https,
        True if job.verify != 0 else False,
        job.ca_file,
        job.start, job.end, job.chunk, job.buf, job.buf_cap,
        &job.done)
    if rc != 0:
        job.status = rc
    return None


# ============================================================
# Public API
# ============================================================

def download(url: str, path: str, threads: INT = 1,
             chunk_size: INT = DL_DEFAULT_CHUNK,
             verify: bool = True, ca_file: str = None) -> INT:
    """Chunked large-file download over http:// or https:// (Range). threads 1=serial,
    2~8=parallel bands; chunk_size bytes (default 256KB). Returns 0/-1/-2/-3/-4 (see last_error()).
    """
    _dl_err_store[0] = 0
    if url is None or path is None or url[0] == 0 or path[0] == 0:
        _dl_set_err("url/path empty")
        return -1
    if chunk_size < 1024:
        chunk_size = DL_DEFAULT_CHUNK
    if threads < 1:
        threads = 1
    if threads > DL_MAX_THREADS:
        threads = DL_MAX_THREADS
    if string.strncmp(url, "http://", 7) != 0 \
            and string.strncmp(url, "https://", 8) != 0:
        _dl_set_err("unsupported scheme")
        return -1

    # Probe total size and Range support
    total: sizet = 0
    prc: INT = _dl_probe(url, verify, ca_file, &total)
    if prc < 0:
        return -3
    if total == 0:
        _dl_set_err("cannot determine size")
        return -3

    # Create (truncate) target file
    if _dl_truncate(path) != 0:
        _dl_set_err("create file failed")
        return -2

    stdio.printf("[download] %s -> %s  total=%zu bytes  threads=%d  chunk=%d\n",
                 url, path, total, threads, chunk_size)

    slab: INT = chunk_size + 8192

    if threads == 1:
        # ---- Serial slices ----
        buf: str = stdlib.malloc(sizet(slab))
        if buf is None:
            _dl_set_err("out of memory")
            return -3
        done: sizet = 0
        rc: INT = _dl_download_band(path, url, 0, verify, ca_file,
                                    0, total, chunk_size, buf, slab, &done)
        stdlib.free(buf)
        if rc != 0:
            return -3 if rc == -1 else -4
        stdio.printf("[download] done: %zu/%zu bytes\n", done, total)
        return 0

    # ---- Multi-thread parallel bands ----
    band: sizet = (total + sizet(threads) - 1) / sizet(threads)
    ctxs: t.CArray[VOIDPTR, DL_MAX_THREADS]
    handles: t.CArray[VOIDPTR, DL_MAX_THREADS]
    threads_ok: INT = 0

    for i in range(int(threads)):
        b_start: sizet = sizet(i) * band
        if b_start >= total:
            break
        b_end: sizet = b_start + band
        if b_end > total:
            b_end = total
        job: DlJob* = stdlib.calloc(1, DlJob.__sizeof__())
        if job is None:
            break
        job.path = _dl_strdup(path)
        job.url = _dl_strdup(url)
        if job.path is None or job.url is None:
            stdlib.free(job.path)
            stdlib.free(job.url)
            stdlib.free(t.CPtr(uint64(job)))
            break
        job.is_https = 1 if (string.strncmp(url, "https://", 8) == 0) else 0
        job.verify = 1 if verify else 0
        job.ca_file = _dl_strdup(ca_file) if ca_file is not None else None
        job.start = b_start
        job.end = b_end
        job.chunk = chunk_size
        job.buf_cap = slab
        job.buf = stdlib.malloc(sizet(slab))
        if job.buf is None:
            stdlib.free(job.path)
            stdlib.free(job.url)
            if job.ca_file is not None:
                stdlib.free(job.ca_file)
            stdlib.free(t.CPtr(uint64(job)))
            break
        job.done = 0
        job.status = 0
        ctxs[i] = t.CPtr(uint64(job))

        th: vt.Thread* = vt.Thread(_dl_worker, ctxs[i])
        if th is None:
            stdlib.free(job.buf)
            stdlib.free(job.path)
            stdlib.free(job.url)
            if job.ca_file is not None:
                stdlib.free(job.ca_file)
            stdlib.free(t.CPtr(uint64(job)))
            break
        th.start()
        handles[i] = th.handle
        threads_ok += 1

    if threads_ok == 0:
        _dl_set_err("thread create failed")
        return -3

    # Wait for all threads to finish
    for j in range(int(threads_ok)):
        wsync.WaitForSingleObject(handles[j], 0xFFFFFFFF)
        wb.CloseHandle(handles[j])

    # Aggregate results
    total_done: sizet = 0
    failed_job: INT = 0
    for k in range(int(threads_ok)):
        jobk: DlJob* = ctxs[k]
        total_done += jobk.done
        if jobk.status != 0:
            failed_job = 1

    # Free thread contexts
    for m in range(int(threads_ok)):
        jobm: DlJob* = ctxs[m]
        stdlib.free(jobm.path)
        stdlib.free(jobm.url)
        if jobm.ca_file is not None:
            stdlib.free(jobm.ca_file)
        stdlib.free(jobm.buf)
        stdlib.free(t.CPtr(uint64(jobm)))

    if failed_job != 0 or total_done != total:
        stdio.printf("[download] incomplete: %zu/%zu bytes\n", total_done, total)
        _dl_set_err("download incomplete")
        return -3
    stdio.printf("[download] done: %zu/%zu bytes\n", total_done, total)
    return 0