#!/usr/bin/env python3
"""List or download Dinox skills. Python 3.10+, no third-party dependencies."""
import argparse
import getpass
import hashlib
import json
import re
import sys
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urlencode, urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener

ORIGIN = "https://aisdk.chatgo.pro"
MAX_ARCHIVE = 70 * 1024 * 1024


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise RuntimeError("Unexpected redirect; request a fresh download URL.")


OPENER = build_opener(NoRedirect())


def api(path, token):
    request = Request(ORIGIN + path, headers={"Authorization": "Bearer " + token})
    with OPENER.open(request, timeout=150) as response:
        raw = response.read(2 * 1024 * 1024 + 1)
    if len(raw) > 2 * 1024 * 1024:
        raise RuntimeError("API response exceeds the example client's limit.")
    result = json.loads(raw)
    if result.get("code") != "000000":
        raise RuntimeError("API returned a non-success response.")
    return result["data"]


def list_skills(token):
    cursor = None
    seen = set()
    while True:
        query = {"page_size": 50}
        if cursor:
            query["start_cursor"] = cursor
        page = api("/api/openapi/skills?" + urlencode(query), token)
        for skill in page["results"]:
            # JSON escapes control characters in user-authored names/descriptions.
            print(json.dumps(skill, ensure_ascii=False))
        if not page["has_more"]:
            return
        cursor = page["next_cursor"]
        if not cursor or cursor in seen:
            raise RuntimeError("Invalid pagination cursor.")
        seen.add(cursor)


def download(info, output):
    url = urlsplit(info["url"])
    if url.scheme != "https" or not url.hostname or url.username or url.password:
        raise RuntimeError("Invalid HTTPS download URL.")
    expected = info["size_bytes"]
    digest = info["sha256"]
    if (info["format"] != "tar.gz" or type(expected) is not int
            or not 0 < expected <= MAX_ARCHIVE
            or not re.fullmatch(r"[a-fA-F0-9]{64}", digest)):
        raise RuntimeError("Invalid archive metadata.")
    # Exclusive creation: never overwrite a user's existing archive.
    with output.open("xb") as target:
        try:
            count = 0
            checksum = hashlib.sha256()
            # Deliberately no Dinox Authorization header on object-storage requests.
            with OPENER.open(Request(info["url"]), timeout=150) as response:
                while chunk := response.read(64 * 1024):
                    count += len(chunk)
                    if count > expected:
                        raise RuntimeError("Archive exceeds the declared size.")
                    checksum.update(chunk)
                    target.write(chunk)
            if count != expected or checksum.hexdigest() != digest.lower():
                raise RuntimeError("Archive size or SHA-256 mismatch.")
        except BaseException:
            target.close()
            output.unlink(missing_ok=True)
            raise
    print(json.dumps({"file": str(output), "id": info["id"],
                      "version_id": info["version_id"], "sha256": digest}, ensure_ascii=False))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--note-id", help="Skill note UUID from the list response")
    parser.add_argument("--output", type=Path, help="New .tar.gz file; never overwritten")
    args = parser.parse_args()
    if bool(args.note_id) != bool(args.output):
        parser.error("--note-id and --output must be used together")
    token = getpass.getpass("Dinox API Token: ").strip()
    if not token:
        parser.error("A token is required")
    if args.note_id:
        info = api("/api/openapi/skills/" + quote(args.note_id, safe=""), token)
        download(info, args.output)
    else:
        list_skills(token)


if __name__ == "__main__":
    try:
        main()
    except HTTPError as error:
        print(f"HTTP {error.code}: see the Skills API error table; refresh expired links.", file=sys.stderr)
        sys.exit(1)
    except (URLError, TimeoutError):
        print("Network request failed; retry later.", file=sys.stderr)
        sys.exit(1)
    except (OSError, ValueError, KeyError, TypeError, RuntimeError):
        # Do not echo exceptions that may include a signed URL or credentials.
        print("Request, file write or validation failed. Check the response and output path.", file=sys.stderr)
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(130)
