from __future__ import annotations
import argparse
import http.cookiejar
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
UA = 'Mozilla/5.0 (cic-downloader)'
LINK_RE = re.compile('href="(browse\\.php\\?p=[^"]+|download\\.php\\?file=[^"]+)"')

def build_opener(cookies_path: Path) -> urllib.request.OpenerDirector:
    jar = http.cookiejar.MozillaCookieJar()
    jar.load(str(cookies_path), ignore_discard=True, ignore_expires=True)
    return urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))

def http_get(opener, url: str, timeout: int=60, retries: int=5) -> bytes:
    last: Exception | None = None
    for attempt in range(retries):
        try:
            req = urllib.request.Request(url, headers={'User-Agent': UA})
            with opener.open(req, timeout=timeout) as resp:
                final = resp.geturl()
                if 'unb.ca/cic/datasets' in final:
                    raise RuntimeError(f'Got redirected to UNB form page ({final}). Token cookie is missing/expired or wrong dataset scope.')
                return resp.read()
        except RuntimeError:
            raise
        except Exception as e:
            last = e
            wait = min(30, 2 ** attempt)
            print(f'  WARN  GET {url} failed ({e!r}); retry in {wait}s ({attempt + 1}/{retries})', file=sys.stderr)
            time.sleep(wait)
    raise RuntimeError(f'GET {url} failed after {retries} attempts: {last!r}')

def list_dir(opener, base: str, p: str) -> list[tuple[str, str]]:
    url = urllib.parse.urljoin(base, 'browse.php') + '?p=' + urllib.parse.quote(p, safe='/')
    html = http_get(opener, url).decode('utf-8', 'replace')
    out: list[tuple[str, str]] = []
    for m in LINK_RE.finditer(html):
        href = m.group(1)
        qs = urllib.parse.parse_qs(urllib.parse.urlparse(href).query)
        if href.startswith('browse.php'):
            out.append(('dir', qs['p'][0]))
        else:
            out.append(('file', qs['file'][0]))
    return out

def walk(opener, base: str, root: str):
    stack = [root]
    seen: set[str] = set()
    while stack:
        p = stack.pop()
        if p in seen:
            continue
        seen.add(p)
        try:
            entries = list_dir(opener, base, p)
        except Exception as e:
            print(f'  WARN  list_dir({p}) failed permanently: {e!r}', file=sys.stderr)
            continue
        for (kind, val) in sorted(entries):
            if kind == 'dir':
                stack.append(val)
            else:
                yield val

def download_file(opener, base: str, remote: str, dest_root: Path, *, root_prefix: str) -> None:
    url = urllib.parse.urljoin(base, 'download.php') + '?file=' + urllib.parse.quote(remote, safe='')
    rel = remote[len(root_prefix):].lstrip('/') if remote.startswith(root_prefix) else remote
    local = dest_root / rel
    local.parent.mkdir(parents=True, exist_ok=True)
    if local.exists() and local.stat().st_size > 0:
        print(f'  SKIP  {rel}  ({local.stat().st_size} bytes, already present)')
        return
    tmp = local.with_suffix(local.suffix + '.part')
    last: Exception | None = None
    for attempt in range(5):
        resume_from = tmp.stat().st_size if tmp.exists() else 0
        try:
            headers = {'User-Agent': UA}
            if resume_from > 0:
                headers['Range'] = f'bytes={resume_from}-'
            req = urllib.request.Request(url, headers=headers)
            t0 = time.monotonic()
            bytes_read = 0
            with opener.open(req, timeout=1800) as resp:
                final = resp.geturl()
                if 'unb.ca/cic/datasets' in final:
                    raise RuntimeError('Token cookie invalid mid-download.')
                status = getattr(resp, 'status', None)
                mode = 'ab'
                if resume_from <= 0:
                    mode = 'wb'
                elif status != 206:
                    print(f'  INFO  {rel} resume request ignored (status={status}); restarting from zero')
                    resume_from = 0
                    mode = 'wb'
                with open(tmp, mode) as fh:
                    while True:
                        buf = resp.read(1 << 20)
                        if not buf:
                            break
                        fh.write(buf)
                        bytes_read += len(buf)
            tmp.replace(local)
            dt = time.monotonic() - t0
            total_bytes = local.stat().st_size
            mb = total_bytes / (1 << 20)
            delta_mb = bytes_read / (1 << 20)
            rate = mb / dt if dt > 0 else 0
            if resume_from > 0:
                resumed_mb = resume_from / (1 << 20)
                rate = delta_mb / dt if dt > 0 else 0
                print(f'  GOT   {rel}  {mb:.1f} MB  +{delta_mb:.1f} MB from {resumed_mb:.1f} MB  {rate:.1f} MB/s')
            else:
                print(f'  GOT   {rel}  {mb:.1f} MB  {rate:.1f} MB/s')
            return
        except urllib.error.HTTPError as e:
            last = e
            if e.code == 416 and resume_from > 0:
                print(f'  WARN  {rel} resume rejected with 416; restarting from zero', file=sys.stderr)
                try:
                    tmp.unlink(missing_ok=True)
                except OSError:
                    pass
                time.sleep(1)
                continue
            wait = min(30, 2 ** attempt)
            print(f'  WARN  {rel} failed ({e!r}); retry in {wait}s ({attempt + 1}/5)', file=sys.stderr)
            time.sleep(wait)
        except RuntimeError:
            raise
        except Exception as e:
            last = e
            wait = min(30, 2 ** attempt)
            print(f'  WARN  {rel} failed ({e!r}); retry in {wait}s ({attempt + 1}/5)', file=sys.stderr)
            time.sleep(wait)
    raise RuntimeError(f'download failed after 5 attempts: {last!r}')

def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--cookies', required=True, type=Path)
    ap.add_argument('--base', required=True, help='dataset URL ending with /, e.g. https://cicresearch.ca/IOTDataset/CIC_IOT_Dataset2023/')
    ap.add_argument('--root', required=True, help='sub-path to crawl (e.g. PCAP or CSV)')
    ap.add_argument('--dest', required=True, type=Path, help='local directory to mirror into')
    ap.add_argument('--dry-run', action='store_true', help='enumerate only; do not download')
    ap.add_argument('--limit', type=int, default=0, help='stop after N files (0 = no limit)')
    ap.add_argument('--skip-ext', default='', help="comma-separated file extensions to skip (e.g. 'zip,7z'); case-insensitive, no dots")
    args = ap.parse_args()
    skip_exts = {e.strip().lower().lstrip('.') for e in args.skip_ext.split(',') if e.strip()}
    if not args.cookies.is_file():
        print(f'ERROR: cookies file not found: {args.cookies}', file=sys.stderr)
        return 2
    opener = build_opener(args.cookies)
    args.dest.mkdir(parents=True, exist_ok=True)
    print(f'Base   : {args.base}')
    print(f'Root   : {args.root}')
    print(f'Dest   : {args.dest}')
    print(f'Walking tree...')
    n_files = 0
    n_skipped = 0
    for remote in walk(opener, args.base, args.root):
        ext = remote.rsplit('.', 1)[-1].lower() if '.' in remote else ''
        if ext in skip_exts:
            n_skipped += 1
            print(f"  SKIP  {remote}  (extension '.{ext}' excluded)")
            continue
        n_files += 1
        if args.dry_run:
            print(f'  FILE  {remote}')
        else:
            try:
                download_file(opener, args.base, remote, args.dest, root_prefix=args.root.rstrip('/'))
            except Exception as e:
                print(f'  FAIL  {remote}: {e}', file=sys.stderr)
        if args.limit and n_files >= args.limit:
            print(f'-- stopped after {args.limit} (--limit) --')
            break
    print(f'Done. {n_files} files processed, {n_skipped} skipped by --skip-ext.')
    return 0
if __name__ == '__main__':
    sys.exit(main())
