#!/usr/bin/env python3
"""Pull a remote Windows folder via the win64_mp file API."""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

AUTH_USER = os.environ.get("AUTH_USER", "admin")
AUTH_PASS = os.environ.get("AUTH_PASS", "blueberries")


def win_join(parent: str, name: str) -> str:
    parent = parent.rstrip("\\/")
    name = name.replace("/", "\\").strip("\\")
    if not name or name in (".", "..") or "\\" in name:
        raise ValueError(f"bad entry name: {name!r}")
    if parent.endswith(":"):
        return parent + "\\" + name
    return parent + "\\" + name


def norm_win(path: str) -> str:
    return path.replace("/", "\\").rstrip("\\").lower()


def should_skip(remote: str, name: str, patterns: list[str]) -> bool:
    remote_n = norm_win(remote)
    name_n = name.replace("/", "\\").strip("\\").lower()
    for raw in patterns:
        pattern = norm_win(raw)
        if not pattern:
            continue
        if "\\" not in pattern and ":" not in pattern:
            if name_n == pattern:
                return True
            continue
        if remote_n == pattern or remote_n.endswith("\\" + pattern) or remote_n.startswith(pattern + "\\"):
            return True
    return False


def split_skips(items: list[str]) -> list[str]:
    out: list[str] = []
    for item in items:
        for part in item.split(","):
            part = part.strip()
            if part:
                out.append(part)
    return out


def prompt(label: str) -> str:
    try:
        value = input(f"{label}: ").strip()
    except EOFError:
        sys.exit("no input")
    if not value:
        sys.exit(f"{label} is required")
    return value


def prompt_optional(label: str) -> str:
    try:
        return input(f"{label}: ").strip()
    except EOFError:
        return ""


def base_url(raw: str) -> str:
    raw = raw.strip().rstrip("/")
    if "://" not in raw:
        raw = "http://" + raw
    parsed = urllib.parse.urlparse(raw)
    if not parsed.hostname:
        sys.exit(f"bad endpoint: {raw}")
    port = parsed.port or 5032
    return f"{parsed.scheme}://{parsed.hostname}:{port}"


def local_root_for(win_path: str) -> Path:
    name = win_path.rstrip("\\/").replace("/", "\\").rsplit("\\", 1)[-1]
    if not name or name.endswith(":"):
        name = "pull"
    dest = Path("pulls") / name
    dest.mkdir(parents=True, exist_ok=True)
    return dest.resolve()


class Agent:
    def __init__(self, origin: str) -> None:
        self.origin = origin
        self.auth = "Basic " + base64.b64encode(f"{AUTH_USER}:{AUTH_PASS}".encode()).decode()

    def request(self, path: str, query: dict[str, str], dest: Path | None = None) -> object | None:
        qs = urllib.parse.urlencode(query)
        url = f"{self.origin}{path}"
        if qs:
            url += "?" + qs
        req = urllib.request.Request(url, headers={"Authorization": self.auth})
        try:
            with urllib.request.urlopen(req, timeout=300) as resp:
                if dest is None:
                    return json.load(resp)
                dest.parent.mkdir(parents=True, exist_ok=True)
                with dest.open("wb") as out:
                    while True:
                        chunk = resp.read(1024 * 1024)
                        if not chunk:
                            break
                        out.write(chunk)
                return None
        except urllib.error.HTTPError as err:
            detail = err.read().decode("utf-8", "replace")
            try:
                detail = json.loads(detail).get("error", detail)
            except json.JSONDecodeError:
                pass
            raise RuntimeError(f"{err.code} {path}: {detail}") from err
        except urllib.error.URLError as err:
            raise RuntimeError(f"connect {self.origin}: {err.reason}") from err

    def health(self) -> None:
        body = self.request("/health", {})
        if not isinstance(body, dict) or body.get("status") != "ok":
            raise RuntimeError(f"unexpected health: {body}")

    def list_dir(self, win_path: str) -> list[dict]:
        body = self.request("/api/v1/files", {"path": win_path, "depth": "0"})
        if not isinstance(body, dict):
            raise RuntimeError(f"bad list response: {body}")
        entries = body.get("entries") or []
        if not isinstance(entries, list):
            raise RuntimeError("list entries is not an array")
        return entries

    def download(self, win_path: str, dest: Path) -> None:
        self.request("/api/v1/download", {"path": win_path}, dest=dest)


def already_have(local: Path, size: object) -> bool:
    if not local.is_file():
        return False
    try:
        return local.stat().st_size == int(size)
    except (TypeError, ValueError):
        return False


def pull(
    agent: Agent, remote: str, dest_root: Path, rel: Path, skips: list[str]
) -> tuple[int, int, int]:
    try:
        entries = agent.list_dir(remote)
    except RuntimeError as err:
        print(f"skip dir {remote}: {err}", file=sys.stderr)
        return 0, 0, 1

    dest_root.joinpath(rel).mkdir(parents=True, exist_ok=True)
    files_ok = 0
    files_skip = 0
    files_fail = 0
    for entry in entries:
        name = str(entry.get("name") or "")
        kind = str(entry.get("type") or "")
        try:
            child = win_join(remote, name)
        except ValueError as err:
            print(f"skip {err}", file=sys.stderr)
            files_fail += 1
            continue
        if should_skip(child, name, skips):
            print(f"skip {child}")
            files_skip += 1
            continue
        if kind == "dir":
            ok, skipped, fail = pull(agent, child, dest_root, rel / name, skips)
            files_ok += ok
            files_skip += skipped
            files_fail += fail
            continue
        if kind != "file":
            print(f"skip {child}: unknown type {kind!r}", file=sys.stderr)
            files_fail += 1
            continue
        local = dest_root / rel / name
        if already_have(local, entry.get("size")):
            files_skip += 1
            continue
        print(f"{child}  ({entry.get('size', '?')} bytes)")
        try:
            agent.download(child, local)
            files_ok += 1
        except RuntimeError as err:
            print(f"  fail: {err}", file=sys.stderr)
            files_fail += 1
    return files_ok, files_skip, files_fail


def self_check() -> None:
    assert win_join(r"C:\Users\me", "docs") == r"C:\Users\me\docs"
    assert win_join(r"C:", "Windows") == r"C:\Windows"
    assert base_url("10.0.0.5:5032") == "http://10.0.0.5:5032"
    assert base_url("http://10.0.0.5") == "http://10.0.0.5:5032"
    cache = r"C:\Users\Larisa\AppData\Local\BraveSoftware\Brave-Browser\User Data\Default\Service Worker\CacheStorage"
    assert should_skip(cache, "CacheStorage", ["CacheStorage"])
    assert should_skip(cache, "CacheStorage", [r"Default\Service Worker\CacheStorage"])
    assert should_skip(cache, "CacheStorage", [cache])
    assert not should_skip(r"C:\Users\Larisa\AppData\Local\BraveSoftware\Brave-Browser\User Data\Default", "Default", ["CacheStorage"])
    assert split_skips(["CacheStorage,GPUCache", "Code Cache"]) == ["CacheStorage", "GPUCache", "Code Cache"]
    print("ok")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Download a remote Windows folder via win64_mp")
    parser.add_argument("endpoint", nargs="?", help="ip:port of the agent")
    parser.add_argument("remote", nargs="?", help="absolute Windows folder path")
    parser.add_argument("dest", nargs="?", help="local destination directory")
    parser.add_argument(
        "--skip",
        "-s",
        action="append",
        default=[],
        metavar="NAME",
        help="folder name or path to ignore; repeatable, comma-separated ok",
    )
    parser.add_argument("--self-check", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.self_check:
        self_check()
        return
    endpoint = args.endpoint or prompt("Endpoint (ip:port)")
    remote = args.remote or prompt("Windows folder path")
    skips = split_skips(args.skip)
    if not skips and args.endpoint is None:
        skips = split_skips([prompt_optional("Skip folders (comma-separated, empty for none)")])
    origin = base_url(endpoint)
    dest = Path(args.dest).resolve() if args.dest else local_root_for(remote)
    dest.mkdir(parents=True, exist_ok=True)

    agent = Agent(origin)
    print(f"checking {origin}/health")
    agent.health()
    print(f"pulling {remote} -> {dest}")
    if skips:
        print("skipping: " + ", ".join(skips))
    ok, skipped, fail = pull(agent, remote, dest, Path(), skips)
    print(f"done: {ok} downloaded, {skipped} skipped, {fail} failures")
    if fail:
        sys.exit(1)


if __name__ == "__main__":
    main()